globalAveragePooling function

Tensor<Vector> globalAveragePooling(Tensor<Matrix> input)

Implementation

Tensor<Vector> globalAveragePooling(Tensor<Matrix> input) {
  int sequenceLength = input.shape[0];
  int dModel = input.shape[1];

  Vector averagedVector = [];
  for (int c = 0; c < dModel; c = c + 1) {
    averagedVector.add(0.0);
  }

  for (int r = 0; r < sequenceLength; r = r + 1) {
    for (int c = 0; c < dModel; c = c + 1) {
      averagedVector[c] = averagedVector[c] + input.data[r * dModel + c];
    }
  }

  for (int c = 0; c < dModel; c = c + 1) {
    averagedVector[c] = averagedVector[c] / sequenceLength;
  }

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

  out.creator = Node(
    [input],
        () {
      for (int r = 0; r < sequenceLength; r = r + 1) {
        for (int c = 0; c < dModel; c = c + 1) {
          int inIdx = r * dModel + c;
          input.grad[inIdx] = input.grad[inIdx] + out.grad[c] / sequenceLength;
        }
      }
    },
    opName: 'global_avg_pool',
    cost: sequenceLength * dModel,
  );

  return out;
}