call method

Tensor call(Tensor x)

Implementation

Tensor call(Tensor x) {
  if (x.shape.length != 4) {
    throw ArgumentError('Conv2d: expected [N, Cin, H, W]; got ${x.shape}');
  }
  final n = x.shape[0];
  final cin = x.shape[1];
  final h = x.shape[2];
  final w = x.shape[3];
  if (cin != inChannels) {
    throw ArgumentError(
      'Conv2d: input channels $cin != declared inChannels $inChannels',
    );
  }
  final hOut = (h + 2 * padding - kernelH) ~/ stride + 1;
  final wOut = (w + 2 * padding - kernelW) ~/ stride + 1;
  if (hOut <= 0 || wOut <= 0) {
    throw ArgumentError(
      'Conv2d: non-positive output size ${[n, outChannels, hOut, wOut]} '
      'for input ${x.shape}, kernel ${[kernelH, kernelW]}, '
      'stride $stride, padding $padding',
    );
  }
  if (x.device != weight.device) {
    // Input can straddle devices: im2col will re-emit `cols` on
    // `weight.device` so the matmul runs there.
  }

  final cols = _im2col(x, n, cin, h, w, hOut, wOut, weight.device);

  // Weight [Cout, Cin, Kh, Kw] -> [Cout, Cin*Kh*Kw] -> transpose to
  // [Cin*Kh*Kw, Cout] for matmul.
  final wFlat = weight.reshape([outChannels, inChannels * kernelH * kernelW]);
  final wT = wFlat.transpose();

  // [N*Hout*Wout, Cin*Kh*Kw] @ [Cin*Kh*Kw, Cout] -> [N*Hout*Wout, Cout].
  var out = cols.matmul(wT);

  if (bias != null) {
    // Broadcast bias [Cout] across every spatial position.
    final bRow = bias!.reshape([1, outChannels]);
    out = out + bRow;
  }

  // Reshape to [N, Hout, Wout, Cout] then permute to [N, Cout, Hout, Wout].
  final nhwc = out.reshape([n, hOut, wOut, outChannels]);
  return _permuteNHWCtoNCHW(nhwc);
}