avgPool2dGPU function

GPUTensor<Matrix> avgPool2dGPU(
  1. GPUTensor<Matrix> input,
  2. int poolSize,
  3. int stride,
  4. CommandBuffer tape, {
  5. GPUTensor<Matrix>? outTensor,
})

Slides a pool of size poolSizexpoolSize over matrix input. The resulting matrix consists of the average of the elements in each window. The value of stride tells how big the step size of the sliding window is. The operations are appended to tape for execution.

Implementation

GPUTensor<Matrix> avgPool2dGPU(
    GPUTensor<Matrix> input, int poolSize, int stride, CommandBuffer tape,
    {GPUTensor<Matrix>? outTensor}) {
  int inHeight = input.shape[0];
  int inWidth = input.shape[1];

  int outHeight = (inHeight - poolSize) ~/ stride + 1;
  int outWidth = (inWidth - poolSize) ~/ stride + 1;

  List<int> outShape = <int>[outHeight, outWidth];
  GPUTensor<Matrix> out = outTensor ?? GPUTensor<Matrix>.empty(outShape);

  tape.putInt(OP_AVG_POOL_2D_FORWARD);
  tape.putString(input.id);
  tape.putString(out.id);
  tape.putInt(poolSize);
  tape.putInt(stride);

  out.creator = GPUNode(
    <GPUTensor>[input],
    (CommandBuffer backwardTape) {
      backwardTape.putInt(OP_AVG_POOL_2D_BACKWARD);
      backwardTape.putString(out.gradId);
      backwardTape.putString(input.gradId);
      backwardTape.putInt(poolSize);
      backwardTape.putInt(stride);
    },
    opName: 'avg_pool_2d_gpu',
  );

  return out;
}