avgPool2dGPU function

GPUTensor<Matrix> avgPool2dGPU(
  1. GPUTensor<Matrix> input,
  2. int poolSize,
  3. int stride,
  4. CommandBuffer tape,
)

Implementation

GPUTensor<Matrix> avgPool2dGPU(
    GPUTensor<Matrix> input,
    int poolSize,
    int stride,
    CommandBuffer tape) {

  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 = 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.id}_grad');
      backwardTape.putString('${input.id}_grad');
      backwardTape.putInt(poolSize);
      backwardTape.putInt(stride);
    },
    opName: 'avg_pool_2d_gpu',
  );

  return out;
}