batchNorm2dGPU function

GPUTensor<Tensor3D> batchNorm2dGPU(GPUTensor<Tensor3D> input, GPUTensor<Vector> gamma, GPUTensor<Vector> beta, GPUTensor<Vector> runningMean, GPUTensor<Vector> runningVariance, double momentum, double epsilon, bool isTraining, CommandBuffer tape, { GPUTensor<Tensor3D>? outTensor, GPUTensor<Vector>? savedMeanTensor, GPUTensor<Vector>? savedInvVarTensor, })

Normalizes a 3D spatial tensor independently for each channel. Uses the mean and variance evaluated across the combined spatial dimensions (height and width) of a given channel. Returns a 3D tensor of identical shape scaled by gamma and shifted by beta. The operations are appended to tape for execution.

Implementation

GPUTensor<Tensor3D> batchNorm2dGPU(
    GPUTensor<Tensor3D> input,
    GPUTensor<Vector> gamma,
    GPUTensor<Vector> beta,
    GPUTensor<Vector> runningMean,
    GPUTensor<Vector> runningVariance,
    double momentum,
    double epsilon,
    bool isTraining,
    CommandBuffer tape,
    {GPUTensor<Tensor3D>? outTensor,
    GPUTensor<Vector>? savedMeanTensor,
    GPUTensor<Vector>? savedInvVarTensor}) {
  int numChannels = input.shape[0];
  int height = input.shape[1];
  int width = input.shape[2];

  GPUTensor<Tensor3D> out =
      outTensor ?? GPUTensor<Tensor3D>.empty(<int>[numChannels, height, width]);
  GPUTensor<Vector> savedMean = savedMeanTensor ??
      GPUTensor<Vector>(List<double>.filled(numChannels, 0.0));
  GPUTensor<Vector> savedInvVar = savedInvVarTensor ??
      GPUTensor<Vector>(List<double>.filled(numChannels, 0.0));
  out.subMap["savedMean"] = savedMean;
  out.subMap["savedInvVar"] = savedInvVar;

  tape.putInt(OP_BATCH_NORM_2D_FORWARD);
  tape.putString(input.id);
  tape.putString(gamma.id);
  tape.putString(beta.id);
  tape.putString(runningMean.id);
  tape.putString(runningVariance.id);
  tape.putString(out.id);
  tape.putString(savedMean.id);
  tape.putString(savedInvVar.id);
  tape.putFloat(momentum);
  tape.putFloat(epsilon);
  tape.putBool(isTraining);

  out.creator = GPUNode(
    [input, gamma, beta],
    (CommandBuffer bTape) {
      bTape.putInt(OP_BATCH_NORM_2D_BACKWARD);
      bTape.putString(out.gradId);
      bTape.putString(input.id);
      bTape.putString(gamma.id);
      bTape.putString(savedMean.id);
      bTape.putString(savedInvVar.id);
      bTape.putString(input.gradId);
      bTape.putString(gamma.gradId);
      bTape.putString(beta.gradId);
    },
    opName: 'batchNorm2dGPU',
    cost: numChannels * height * width * 4,
  );

  return out;
}