conv2d function

Tensor<Matrix> conv2d(
  1. Tensor<Matrix> input,
  2. Tensor<Matrix> kernel, {
  3. String padding = 'valid',
})

Implementation

Tensor<Matrix> conv2d(
  Tensor<Matrix> input,
  Tensor<Matrix> kernel, {
  String padding = 'valid',
}) {
  Matrix inputMatrix = input.value;
  int padSize = 0;
  if (padding == 'same') {
    padSize = (kernel.value.length - 1) ~/ 2;
    inputMatrix = padMatrix(input.value, padSize);
  }
  int inputHeight = inputMatrix.length;
  int inputWidth = inputMatrix[0].length;
  int kernelHeight = kernel.value.length;
  int kernelWidth = kernel.value[0].length;
  int outputHeight = inputHeight - kernelHeight + 1;
  int outputWidth = inputWidth - kernelWidth + 1;

  Matrix outputValue = [];
  for (int i = 0; i < outputHeight; i++) {
    Vector row = [];
    for (int j = 0; j < outputWidth; j++) {
      row.add(0.0);
    }
    outputValue.add(row);
  }

  for (int y = 0; y < outputHeight; y++) {
    for (int x = 0; x < outputWidth; x++) {
      double sum = 0;
      for (int ky = 0; ky < kernelHeight; ky++) {
        for (int kx = 0; kx < kernelWidth; kx++) {
          sum += inputMatrix[y + ky][x + kx] * kernel.value[ky][kx];
        }
      }
      outputValue[y][x] = sum;
    }
  }
  Tensor<Matrix> out = Tensor<Matrix>(outputValue);
  int cost = outputHeight * outputWidth * 2 * kernelHeight * kernelWidth;
  out.creator = Node(
    [input, kernel],
    () {
      for (int y = 0; y < outputHeight; y++) {
        for (int x = 0; x < outputWidth; x++) {
          for (int ky = 0; ky < kernelHeight; ky++) {
            for (int kx = 0; kx < kernelWidth; kx++) {
              if (padding == 'same' &&
                  (y + ky < padSize ||
                      y + ky >= input.value.length + padSize ||
                      x + kx < padSize ||
                      x + kx >= input.value[0].length + padSize))
                continue;
              int inputGradY = (padding == 'same') ? y + ky - padSize : y + ky;
              int inputGradX = (padding == 'same') ? x + kx - padSize : x + kx;
              input.grad[inputGradY][inputGradX] +=
                  kernel.value[ky][kx] * out.grad[y][x];
              kernel.grad[ky][kx] +=
                  inputMatrix[y + ky][x + kx] * out.grad[y][x];
            }
          }
        }
      }
    },
    opName: 'conv2d',
    extraParams: {
      'padding': padding,
    }, // <-- CRITICAL: Storing the non-Tensor parameter
    cost: cost,
  );
  return out;
}