avgPool1d function

Tensor<Vector> avgPool1d(
  1. Tensor<Vector> input,
  2. int poolSize,
  3. int stride
)

Implementation

Tensor<Vector> avgPool1d(Tensor<Vector> input, int poolSize, int stride) {
  int inputSize = input.shape[0];
  int outputSize = (inputSize - poolSize) ~/ stride + 1;

  Vector outputValue = [];

  for (int i = 0; i < outputSize; i = i + 1) {
    double sum = 0.0;
    for (int p = 0; p < poolSize; p = p + 1) {
      sum = sum + input.data[i * stride + p];
    }
    outputValue.add(sum / poolSize);
  }

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

  out.creator = Node(
    [input],
        () {
      double gradDist = 1.0 / poolSize;
      for (int i = 0; i < outputSize; i = i + 1) {
        for (int p = 0; p < poolSize; p = p + 1) {
          int inIdx = i * stride + p;
          input.grad[inIdx] = input.grad[inIdx] + out.grad[i] * gradDist;
        }
      }
    },
    opName: 'avg_pool_1d',
    cost: outputSize * poolSize,
  );

  return out;
}