maxPool2dGPU function

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

Implementation

GPUTensor<Matrix> maxPool2dGPU(GPUTensor<Matrix> input, int poolSize, int stride, CommandBuffer tape) {
  int inputHeight = input.shape[0];
  int inputWidth = input.shape[1];
  int outputHeight = (inputHeight - poolSize) ~/ stride + 1;
  int outputWidth = (inputWidth - poolSize) ~/ stride + 1;

  GPUTensor<Matrix> out = GPUTensor<Matrix>.empty(<int>[outputHeight, outputWidth]);
  GPUTensor<Matrix> indices = GPUTensor<Matrix>.empty(<int>[outputHeight, outputWidth]);

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

  out.creator = GPUNode(
    [input],
        (CommandBuffer bTape) {
      bTape.putInt(OP_MAX_POOL_2D_BACKWARD);
      bTape.putString('${out.id}_grad');
      bTape.putString(indices.id);
      bTape.putString('${input.id}_grad');
    },
    opName: 'maxPool2dGPU',
    cost: inputHeight * inputWidth,
  );

  return out;
}