Convolutional networks

Deep dive

How it works

In: [N, C_in, H, W] input, weights, bias · Out: [N, C_out, H_out, W_out] · State: weights, BatchNorm running stats · Cost: 2·C_out·(C_in/g)·k²·H_out·W_out FLOPs.

1. Slide the k×k window over the padded input; dot it with each filter; add the bias.
2. Normalize, apply the nonlinearity; stride shrinks the grid.

Tensors in NCHW (or NHWC): input [N, C_in, H, W], weight [C_out, C_in/g, k, k], output [N, C_out, H_out, W_out].

y[n, co, oy, ox] = b[co] + Σ_{ci in group(co)} Σ_{ky,kx} x[n, ci, oy·s + ky·d − p, ox·s + kx·d − p] · w[co, ci, ky, kx]

• Stride s downsamples; padding p keeps borders (zeros outside); dilation d spreads taps to widen the receptive field with no extra cost.
• Groups g split channels into g independent convs. g = C_in = C_out is depthwise (one filter per channel); a following 1×1 conv (pointwise) mixes channels.
• Receptive field: with jump j₀ = 1, r₀ = 1: r_l = r_{l−1} + (k_l − 1)·j_{l−1}, j_l = j_{l−1}·s_l. Two 3×3 stride-1 layers see 5×5; three see 7×7 with fewer params than one 7×7 (27C² vs 49C²).
• Block patterns: ResNet basic block (3×3, 3×3, identity add); bottleneck (1×1 reduce to C/4, 3×3, 1×1 expand, add); MobileNetV2 inverted residual (1×1 expand ×6, 3×3 depthwise, 1×1 project, linear output); ConvNeXt (7×7 depthwise, LayerNorm, 1×1 to 4C, GELU, 1×1 back).
• Normalization: conv → BatchNorm → ReLU. BN normalizes each channel over (N, H, W). At inference it is an affine map and folds into the conv: w′ = w·γ/√(σ²+ε), b′ = (b − μ)·γ/√(σ²+ε) + β.
• Lowering to GEMM: im2col builds a [C_in·k², H_out·W_out] matrix of patches; W reshaped to [C_out, C_in·k²] times it gives the output. cuDNN's implicit GEMM does the same gather on the fly in shared memory, without materializing the k²-times-larger matrix.

Code

Direct convolution with stride, padding, dilation and groups; im2col + GEMM for comparison; BatchNorm folding. The checks compare both paths and verify folding.

// 2D convolution on CHW tensors (one image): direct loop and im2col + GEMM, plus BatchNorm folding.
#[derive(Clone, Copy)]
pub struct Conv { pub c_in: usize, pub c_out: usize, pub k: usize, pub stride: usize, pub pad: usize, pub dil: usize, pub groups: usize }

impl Conv {
    pub fn out_size(&self, n: usize) -> usize {
        (n + 2 * self.pad - self.dil * (self.k - 1) - 1) / self.stride + 1
    }
    // Input coordinate for output o and kernel tap kk; None = inside the zero padding.
    fn src(&self, o: usize, kk: usize, n: usize) -> Option<usize> {
        let i = (o * self.stride + kk * self.dil) as isize - self.pad as isize;
        if i >= 0 && (i as usize) < n { Some(i as usize) } else { None }
    }
    /// x: [c_in, h, w]; wt: [c_out, c_in/groups, k, k]; bias: [c_out] -> [c_out, ho, wo]
    pub fn forward(&self, x: &[f32], h: usize, w: usize, wt: &[f32], bias: &[f32]) -> Vec<f32> {
        let (ho, wo, k) = (self.out_size(h), self.out_size(w), self.k);
        let (cig, cog) = (self.c_in / self.groups, self.c_out / self.groups);
        let mut y = vec![0.0; self.c_out * ho * wo];
        for co in 0..self.c_out {
            let g = co / cog; // group of this output channel sees input channels g*cig..(g+1)*cig
            for oy in 0..ho {
                for ox in 0..wo {
                    let mut acc = bias[co];
                    for ci in 0..cig {
                        let plane = &x[(g * cig + ci) * h * w..][..h * w];
                        for ky in 0..k {
                            let Some(iy) = self.src(oy, ky, h) else { continue };
                            for kx in 0..k {
                                let Some(ix) = self.src(ox, kx, w) else { continue };
                                acc += plane[iy * w + ix] * wt[((co * cig + ci) * k + ky) * k + kx];
                            }
                        }
                    }
                    y[(co * ho + oy) * wo + ox] = acc;
                }
            }
        }
        y
    }
    /// groups = 1: unfold patches into cols [c_in*k*k, ho*wo], then y = W[c_out, c_in*k*k] · cols.
    pub fn forward_im2col(&self, x: &[f32], h: usize, w: usize, wt: &[f32], bias: &[f32]) -> Vec<f32> {
        let (ho, wo, k) = (self.out_size(h), self.out_size(w), self.k);
        let (rows, cols_n) = (self.c_in * k * k, ho * wo);
        let mut cols = vec![0.0; rows * cols_n];
        for ci in 0..self.c_in { for ky in 0..k { for kx in 0..k {
            let r = (ci * k + ky) * k + kx;
            for oy in 0..ho { for ox in 0..wo {
                if let (Some(iy), Some(ix)) = (self.src(oy, ky, h), self.src(ox, kx, w)) {
                    cols[r * cols_n + oy * wo + ox] = x[(ci * h + iy) * w + ix];
                }
            }}
        }}}
        let mut y = vec![0.0; self.c_out * cols_n];
        for co in 0..self.c_out {
            let yr = &mut y[co * cols_n..][..cols_n];
            yr.fill(bias[co]);
            for r in 0..rows {
                let a = wt[co * rows + r];
                for (yv, cv) in yr.iter_mut().zip(&cols[r * cols_n..][..cols_n]) { *yv += a * cv; }
            }
        }
        y
    }
}

/// Fold inference-mode BatchNorm (gamma, beta, running mean/var) into the preceding conv.
pub fn fold_bn(wt: &mut [f32], bias: &mut [f32], gamma: &[f32], beta: &[f32], mean: &[f32], var: &[f32], eps: f32) {
    let per = wt.len() / bias.len();
    for co in 0..bias.len() {
        let s = gamma[co] / (var[co] + eps).sqrt();
        wt[co * per..(co + 1) * per].iter_mut().for_each(|w| *w *= s);
        bias[co] = (bias[co] - mean[co]) * s + beta[co];
    }
}

fn main() {
    let mut seed = 7u64;
    let mut rnd = |n: usize| -> Vec<f32> { (0..n).map(|_| {
        seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; (seed % 2001) as f32 / 1000.0 - 1.0 }).collect() };
    let c = Conv { c_in: 3, c_out: 8, k: 3, stride: 2, pad: 1, dil: 1, groups: 1 };
    let (h, w) = (9, 7);
    let (x, wt, b) = (rnd(3 * h * w), rnd(8 * 3 * 9), rnd(8));
    let (y1, y2) = (c.forward(&x, h, w, &wt, &b), c.forward_im2col(&x, h, w, &wt, &b));
    assert_eq!((c.out_size(h), c.out_size(w)), (5, 4));
    assert!(y1.iter().zip(&y2).all(|(a, b)| (a - b).abs() < 1e-4));
    // Depthwise 3x3 (groups = channels) followed by BN folding gives the same output as conv then BN.
    let dw = Conv { c_in: 3, c_out: 3, k: 3, stride: 1, pad: 1, dil: 1, groups: 3 };
    let (mut wd, mut bd) = (rnd(3 * 9), vec![0.0; 3]);
    let (g, be, mu, var) = (vec![1.5, 0.5, 2.0], vec![0.1, -0.2, 0.3], vec![0.2, 0.0, -0.1], vec![4.0, 1.0, 0.25]);
    let raw = dw.forward(&x, h, w, &wd, &bd);
    fold_bn(&mut wd, &mut bd, &g, &be, &mu, &var, 1e-5);
    let fused = dw.forward(&x, h, w, &wd, &bd);
    for (i, (r, f)) in raw.iter().zip(&fused).enumerate() {
        let c = i / (h * w);
        let bn = (r - mu[c]) / (var[c] + 1e-5f32).sqrt() * g[c] + be[c];
        assert!((bn - f).abs() < 1e-4);
    }
    println!("ok");
}

Complexity and constants

• Params: C_out·(C_in/g)·k² (+ C_out bias, usually dropped before BN).
• FLOPs: 2·C_out·(C_in/g)·k²·H_out·W_out. Halving resolution cuts FLOPs 4×; ResNets double channels at each stride-2 stage so FLOPs per stage stay roughly constant.
• Arithmetic intensity: a dense 3×3 conv reuses each input value C_out·9 times and is compute-bound on GPUs; depthwise reuses it only 9 times and is memory-bound. MobileNet-style FLOP savings show up on CPUs and NPUs more than on big GPUs.
• im2col memory: k² × the input activation (9× for 3×3), which is why production kernels use implicit GEMM.
• Winograd F(2×2, 3×3) uses 16 multiplies per 2×2 output tile instead of 36 (2.25×); F(4×4, 3×3) 4×, at some loss of fp16 accuracy. FFT convolution wins only for large kernels.

Efficiency

• Cost: 231M FLOPs for a 3×3 conv, 64→64 channels at 56×56; ResNet-50 about 4.1 G multiply-adds per image. Dense convs are compute-bound on GPUs, depthwise memory-bound.
• Levers:
  • depthwise-separable conv — about 8-9x fewer FLOPs for 3×3 (MobileNets) — edge
  • Winograd F(2×2, 3×3) — 2.25x fewer multiplies (Lavin and Gray) — 3×3 stride-1
  • BatchNorm folding, conv+ReLU fusion — one kernel — inference
  • channels-last, channels multiple of 8 — tensor-core tiles (Implementation notes) — mixed precision
  • implicit GEMM instead of im2col — no k²-times-larger buffer (cuDNN paper) — GPU kernels
• Naive → standard → state of the art:
  • direct loops · far from peak
  • im2col + GEMM · k² times the activation memory
  • implicit GEMM or Winograd, fused, channels-last · compute-bound

Implementation notes

• Layout: tensor-core conv kernels work in NHWC (channels-last). With NCHW tensors the framework inserts transposes; in PyTorch use memory_format=torch.channels_last for both model and input under mixed precision.
• Channel counts as multiples of 8 (fp16/bf16) or 16 (int8/fp8) so channels map onto tensor-core tiles; pad the 3-channel RGB stem to 4 or 8.
• Init: He normal, std = √(2/fan_in) with fan_in = (C_in/g)·k², for ReLU nets. Zero-init the last BN γ in each residual branch so every block starts as identity (Goyal et al.).
• BatchNorm: PyTorch defaults ε = 1e-5, momentum 0.1. Below about 16 images per device the batch statistics get noisy: use SyncBN, or GroupNorm with 32 groups.
• Fuse at inference: conv + folded BN + ReLU (+ residual add) is one kernel; see kernel fusion.
• Padding semantics: PyTorch pads symmetrically and padding='same' only supports stride 1. TensorFlow SAME with stride 2 puts the extra pad on the bottom/right, so ported weights shift by one pixel unless you replicate it.

Pitfalls

• Forgetting eval mode: BN keeps using batch statistics and updating running stats at inference.
• Weight decay on BN γ/β and biases: usually hurts; exclude them.
• Strided conv and max-pool alias: outputs change when the input shifts by one pixel. Blur before subsampling (Zhang 2019) if shift-invariance matters.
• Running BN statistics computed on augmented train data, then serving at a different resolution: recalibrate the statistics at the test resolution.
• usize underflow in the output-size formula when the kernel (with dilation) is larger than the padded input: validate shapes up front.
• Measuring speed in FLOPs: depthwise and grouped convs often run slower than their FLOP count suggests.

Variants and alternatives

In production

cuDNN, oneDNN and XNNPACK conv kernels; torchvision ResNet/EfficientNet/ConvNeXt; YOLO-family detectors; the U-Nets in Stable Diffusion 1.x/2.x and SDXL; WaveNet for speech synthesis; U-Net segmentation of biomedical images; the residual towers of AlphaZero-style game engines.

Implementations

Sources

• LeCun, Bottou, Bengio, Haffner, Gradient-Based Learning Applied to Document Recognition, Proc. IEEE 1998. paper
• He, Zhang, Ren, Sun, Deep Residual Learning for Image Recognition, CVPR 2016. paper
• He, Zhang, Ren, Sun, Delving Deep into Rectifiers, ICCV 2015. paper
• Ioffe and Szegedy, Batch Normalization, ICML 2015. paper
• Howard et al., MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications, 2017. paper
• Lavin and Gray, Fast Algorithms for Convolutional Neural Networks, CVPR 2016. paper
• Chetlur et al., cuDNN: Efficient Primitives for Deep Learning, 2014. paper
• Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, 2017. paper
• Zhang, Making Convolutional Networks Shift-Invariant Again, ICML 2019. paper
• Liu et al., A ConvNet for the 2020s, CVPR 2022. paper
• Dumoulin and Visin, A guide to convolution arithmetic for deep learning, 2016. paper
• van den Oord et al., WaveNet: A Generative Model for Raw Audio, 2016. paper
• Ronneberger, Fischer, Brox, U-Net: Convolutional Networks for Biomedical Image Segmentation, MICCAI 2015. paper