maxPool2d function

Tensor<Matrix> maxPool2d(Tensor<Matrix> input, int poolSize, int stride)

Implementation

Tensor<Matrix> maxPool2d(Tensor<Matrix> input, int poolSize, int stride) {
  int inputHeight = input.shape[0];
  int inputWidth = input.shape[1];

  int outputHeight = (inputHeight - poolSize) ~/ stride + 1;
  int outputWidth = (inputWidth - poolSize) ~/ stride + 1;

  Matrix outputValue = [];
  List<int> maxIndices = [];

  for (int y = 0; y < outputHeight; y = y + 1) {
    Vector row = [];
    for (int x = 0; x < outputWidth; x = x + 1) {
      double maxVal = -double.infinity;
      int maxIdx = -1;
      for (int py = 0; py < poolSize; py = py + 1) {
        int inY = y * stride + py;
        for (int px = 0; px < poolSize; px = px + 1) {
          int inX = x * stride + px;
          int flatInIdx = inY * inputWidth + inX;
          double val = input.data[flatInIdx];
          if (val > maxVal) {
            maxVal = val;
            maxIdx = flatInIdx;
          }
        }
      }
      row.add(maxVal);
      maxIndices.add(maxIdx);
    }
    outputValue.add(row);
  }

  Tensor<Matrix> out = Tensor<Matrix>(outputValue);

  out.creator = Node(
    [input],
        () {
      int outLength = out.data.length;
      for (int i = 0; i < outLength; i = i + 1) {
        int mIdx = maxIndices[i];
        input.grad[mIdx] = input.grad[mIdx] + out.grad[i];
      }
    },
    opName: 'max_pool_2d',
    cost: outputHeight * outputWidth * poolSize * poolSize,
  );

  return out;
}