avgPool2d function

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

Implementation

Tensor<Matrix> avgPool2d(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 = [];
  double numElements = (poolSize * poolSize).toDouble();

  for (int y = 0; y < outputHeight; y = y + 1) {
    Vector row = [];
    for (int x = 0; x < outputWidth; x = x + 1) {
      double sum = 0.0;
      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;
          sum = sum + input.data[inY * inputWidth + inX];
        }
      }
      row.add(sum / numElements);
    }
    outputValue.add(row);
  }

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

  out.creator = Node(
    [input],
        () {
      double gradDist = 1.0 / numElements;
      for (int y = 0; y < outputHeight; y = y + 1) {
        for (int x = 0; x < outputWidth; x = x + 1) {
          int outIdx = y * outputWidth + x;
          for (int py = 0; py < poolSize; py = py + 1) {
            for (int px = 0; px < poolSize; px = px + 1) {
              int inY = y * stride + py;
              int inX = x * stride + px;
              int inIdx = inY * inputWidth + inX;
              input.grad[inIdx] = input.grad[inIdx] + out.grad[outIdx] * gradDist;
            }
          }
        }
      }
    },
    opName: 'avg_pool_2d',
    cost: outputHeight * outputWidth * poolSize * poolSize,
  );

  return out;
}