Implement Conv2D from Scratch

Math Foundations

Given input image tensor $X \in \mathbb{R}^{B \times C_{\text{in}} \times H \times W}$ and learnable filter kernels $W \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times K_H \times K_W}$:

A 2D Convolution slides each filter across spatial dimensions with stride $S$ and padding $P$, computing element-wise dot products:

$$Y_{b, c_{\text{out}}, i, j} = b_{c_{\text{out}}} + \sum_{c_{\text{in}}=1}^{C_{\text{in}}} \sum_{m=1}^{K_H} \sum_{n=1}^{K_W} X_{b, c_{\text{in}}, i \cdot S + m, j \cdot S + n} \cdot W_{c_{\text{out}}, c_{\text{in}}, m, n}$$

Output spatial dimensions $H_{\text{out}}$ and $W_{\text{out}}$:

$$H_{\text{out}} = \left\lfloor \frac{H + 2P - K_H}{S} \right\rfloor + 1, \quad W_{\text{out}} = \left\lfloor \frac{W + 2P - K_W}{S} \right\rfloor + 1$$

Input Image (H x W) ──► Slide Filter Kernel (Kh x Kw) ──► Compute Dot Product ──► Output Feature Map (Hout x Wout)

Naive Loops vs im2col Vectorization

A naive implementation requires 6 nested Python for loops (Batch, Out Channels, In Channels, Out Height, Out Width, Kernel Height, Kernel Width), taking seconds per image.

The production im2col algorithm:

1. Extracts all $K_H \times K_W \times C_{\text{in}}$ receptive field patches into a 2D matrix of shape [K_H * K_W * C_in, H_out * W_out].
2. Reshapes filters to [C_out, K_H * K_W * C_in].
3. Performs a single GEMM (General Matrix Multiplication): Filters @ Patches!

NumPy Implementation using Spatial Sliding Windows

import numpy as np


class Conv2DFromScratch:

  def __init__(
      self, in_channels, out_channels, kernel_size, stride=1, padding=0
  ):
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.kernel_size = kernel_size
    self.stride = stride
    self.padding = padding

    # He Initialization for weights
    scale = np.sqrt(2.0 / (in_channels * kernel_size * kernel_size))
    self.W = (
        np.random.randn(out_channels, in_channels, kernel_size, kernel_size)
        * scale
    )
    self.b = np.zeros((out_channels, 1, 1))

  def forward(self, X):
    # X shape: [B, C_in, H, W]
    B, C, H, W = X.shape
    K = self.kernel_size
    S = self.stride
    P = self.padding

    # Apply Zero Padding
    if P > 0:
      X_padded = np.pad(
          X, ((0, 0), (0, 0), (P, P), (P, P)), mode='constant', constant_values=0
      )
    else:
      X_padded = X

    H_padded, W_padded = X_padded.shape[2], X_padded.shape[3]
    H_out = (H_padded - K) // S + 1
    W_out = (W_padded - K) // S + 1

    Y = np.zeros((B, self.out_channels, H_out, W_out))

    # Vectorized spatial sliding window loop
    for i in range(H_out):
      for j in range(W_out):
        h_start = i * S
        h_end = h_start + K
        w_start = j * S
        w_end = w_start + K

        # Extract 3D slice across all batch samples: [B, C_in, K, K]
        slice_x = X_padded[:, :, h_start:h_end, w_start:w_end]

        # Compute dot product over all filters
        for c in range(self.out_channels):
          # Element-wise multiply and sum over in_channels, H, W
          Y[:, c, i, j] = (
              np.sum(slice_x * self.W[c, :, :, :], axis=(1, 2, 3)) + self.b[c]
          )

    return Y

Say this out loud

2D Convolutions slide filter kernels across spatial image dimensions to extract local visual features. Spatial padding preserves boundary dimensions, while stride controls downsampling. Production implementations use the im2col transformation to convert 3D spatial sliding patches into 2D matrix multiplications for fast parallel GPU execution.

Followups to expect

1. What is Depthwise Separable Convolution? Splitting a 2D convolution into a spatial Depthwise convolution per channel followed by a 1x1 Pointwise channel mixing convolution, reducing computational FLOPs by 8 to 9 times in MobileNet.
2. What is Transposed Convolution (Deconvolution)? An upside down convolution operation that upsamples spatial dimensions, commonly used in image segmentation and Generative Adversarial Networks.