globalAveragePooling function

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

Computes the average of a Tensor

Implementation

Tensor<Vector> globalAveragePooling(Tensor<Matrix> input) {
  Matrix inputMatrix = input.value as Matrix;
  int sequenceLength = inputMatrix.length;
  int dModel = inputMatrix[0].length;

  Vector averagedVector = List<double>.filled(dModel, 0.0);

  // Summing elements across the sequence dimension (rows)
  for (int r = 0; r < sequenceLength; r++) {
    for (int c = 0; c < dModel; c++) {
      averagedVector[c] += inputMatrix[r][c];
    }
  }

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

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

  // Backward pass: Gradient is distributed equally back to all rows of the input matrix
  out.creator = Node(<Tensor>[input], () {
    for (int r = 0; r < sequenceLength; r++) {
      for (int c = 0; c < dModel; c++) {
        input.grad[r][c] += out.grad[c] / sequenceLength.toDouble();
      }
    }
  }, opName: 'global_avg_pool', cost: sequenceLength * dModel);

  return out;
}