How to Implement Convolutional Neural Networks from Scratch: A Complete Guide to TinyTorch

TLDR: You can implement convolutional neural networks from scratch by combining explicit nested loops for the convolution operation, automatic gradient tracking through a custom autograd engine, and modular spatial layers like Conv2d and MaxPool2d, as demonstrated in the TinyTorch educational framework from Harvard's cs249r_book repository.

Implementing convolutional neural networks from scratch requires understanding both the mathematical operations and the computational graph mechanics that enable learning. The TinyTorch library in the harvard-edge/cs249r_book repository provides a pedagogical implementation that mirrors textbook definitions while maintaining readable, educational code. This guide walks through the exact implementation details found in the source files, from the seven-nested-loop convolution kernel to the backward pass gradient propagation.

Core Components of the TinyTorch CNN Implementation

Spatial Layers and Input Validation

The convolutional and pooling operations live in tinytorch/src/09_convolutions/09_convolutions.py. This module defines Conv2d, MaxPool2d, and AvgPool2d with strict input validation through validate_4d_input, ensuring all inputs conform to the expected (batch, channel, height, width) tensor format.

The Tensor and Autograd System

The automatic differentiation engine resides in tinytorch/core/tensor.py and tinytorch/core/autograd.py. The Tensor class wraps NumPy arrays and stores gradient information, while the Function base class in autograd.py defines the interface for forward and backward operations. When enable_autograd() is active, every operation builds a computational graph node via _grad_fn attributes.

The Six-Step Convolution Pipeline

According to the TinyTorch source code, the Conv2d.forward method implements a strict six-step pipeline:

1. Input Validation – The validate_4d_input function checks that the input x is a 4-D tensor with shape (batch, channel, height, width).

2. Output Shape Computation – The method _compute_output_shape calculates spatial dimensions using the standard formula ((in + 2·padding – kernel) // stride) + 1.

3. Zero-Padding – When padding > 0, the _apply_padding method adds symmetric zeros around the spatial dimensions of the input data.

4. Sliding-Window Convolution – The _convolve_loops method executes explicit seven-nested loops to perform multiply-accumulate operations for every output pixel, iterating over batch, output channels, input channels, and spatial dimensions.

5. Bias Addition – If the layer includes bias parameters, the forward pass adds a learnable bias term per output channel using out[:, oc, :, :] += self.bias.data[oc].

6. Autograd Registration – The result is wrapped in a Tensor with requires_grad flags set, and a Conv2dBackward function is assigned to _grad_fn for gradient propagation.

Forward Pass Implementation Details

The forward pass in tinytorch/src/09_convolutions/09_convolutions.py explicitly implements the convolution operation without vectorization libraries:

def forward(self, x):
    # 1️⃣ Validate input shape

    validate_4d_input(x, "Conv2D")

    batch, _, h, w = x.shape
    # 2️⃣ Output spatial dimensions

    out_h, out_w = self._compute_output_shape(h, w)
    # 3️⃣ Pad if needed

    padded = self._apply_padding(x.data)
    # 4️⃣ Perform the explicit loops

    out = self._convolve_loops(padded, batch, out_h, out_w)

    # 5️⃣ Add bias (if present)

    if self.bias is not None:
        for oc in range(self.out_channels):
            out[:, oc, :, :] += self.bias.data[oc]

    # 6️⃣ Return a Tensor with gradient tracking

    result = Tensor(out,
                    requires_grad=(x.requires_grad or self.weight.requires_grad))
    if result.requires_grad:
        result._grad_fn = Conv2dBackward(
            x, self.weight, self.bias,
            self.stride, self.padding, self.kernel_size,
            padded.shape)
    return result

Backward Pass and Gradient Flow

The Conv2dBackward class in the same file mirrors the forward loops to distribute the upstream gradient. It pads the input gradient when the forward pass used padding, then removes the padding before returning the gradient with respect to the original input. This ensures that gradients flow correctly through the computational graph to update both the filter weights and the bias terms during backpropagation.

Pooling Layers Implementation

Both MaxPool2d and AvgPool2d follow the same three-step pattern of validation, padding, and sliding windows. The MaxPool2dBackward class stores the positions of the maxima (indices) to route gradients only to the winning neurons, while AvgPool2dBackward evenly distributes the gradient across all positions in the pooling window.

Weight Initialization Strategy

TinyTorch uses He (Kaiming) initialization for convolutional layers, which is critical for training deep networks with ReLU activations. The implementation in Conv2d.__init__ calculates the fan-in as in_channels * kernel_h * kernel_w and initializes weights from a normal distribution with standard deviation sqrt(2.0 / fan_in):

fan_in = in_channels * kernel_h * kernel_w
std = np.sqrt(2.0 / fan_in)
self.weight = Tensor(np.random.normal(0, std,
                (out_channels, in_channels, kernel_h, kernel_w)),
                requires_grad=True)

Complete CNN Model Example

Here is a runnable example from the repository demonstrating how to wire these components into a complete network for MNIST-style data:

import numpy as np
from tinytorch.core.tensor import Tensor
from tinytorch.core.autograd import enable_autograd
from tinytorch.src.09_convolutions.09_convolutions import Conv2d, MaxPool2d, AvgPool2d
from tinytorch.src.04_linear.04_linear import Linear
from tinytorch.src.06_activation.06_activation import ReLU, LogSoftmax

enable_autograd()

class SimpleCNN:
    def __init__(self):
        self.conv1 = Conv2d(in_channels=1, out_channels=8, kernel_size=3, padding=1)
        self.pool1 = MaxPool2d(kernel_size=2, stride=2)   # → (14×14)

        self.conv2 = Conv2d(in_channels=8, out_channels=16, kernel_size=3, padding=1)
        self.pool2 = AvgPool2d(kernel_size=2, stride=2)   # → (7×7)

        self.fc = Linear(in_features=16*7*7, out_features=10)

        self.relu = ReLU()
        self.logsoftmax = LogSoftmax()

    def __call__(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool1(x)
        x = self.relu(self.conv2(x))
        x = self.pool2(x)

        # Flatten: (batch, C, H, W) → (batch, C*H*W)

        batch = x.shape[0]
        x = Tensor(x.data.reshape(batch, -1))

        x = self.fc(x)
        return self.logsoftmax(x)

# Dummy data: 32‑pixel grayscale images, batch size 4

inputs = Tensor(np.random.randn(4, 1, 32, 32))
targets = Tensor(np.eye(10)[np.random.choice(10, 4)])   # one‑hot labels

model = SimpleCNN()
outputs = model(inputs)

# Simple cross‑entropy loss (negative log‑likelihood)

loss = -(targets * outputs).sum() / outputs.shape[0]
loss.backward()

print(f"Loss: {loss.item():.4f}")

# Gradients are now stored in each parameter's `.grad` attribute

This example demonstrates the layer composition pattern: Conv2d → ReLU → MaxPool2d → Conv2d → AvgPool2d → Flatten → Linear → LogSoftmax. Calling loss.backward() triggers the chain of *_Backward.apply functions across all layers, storing gradients in model.conv1.weight.grad and other parameter attributes.

Key Source Files Reference

To explore the implementation further, examine these specific files in the harvard-edge/cs249r_book repository:

• tinytorch/src/09_convolutions/09_convolutions.py – Contains Conv2d, MaxPool2d, AvgPool2d, validation helpers, and the forward/backward logic.
• tinytorch/core/tensor.py – Implements the Tensor class with shape handling and gradient storage.
• tinytorch/core/autograd.py – Defines the Function base class and enable_autograd() context.
• tinytorch/src/04_linear/04_linear.py – Provides the Linear layer for fully-connected operations.
• tinytorch/src/06_activation/06_activation.py – Houses ReLU, LogSoftmax, and other pointwise activations.

Summary

• Implement convolutional neural networks from scratch using explicit nested loops in _convolve_loops to handle the sliding-window multiply-accumulate operations.
• Validate all inputs with validate_4d_input to enforce (batch, channel, height, width) tensor shapes.
• Compute output dimensions using the standard formula ((in + 2·padding – kernel) // stride) + 1 in _compute_output_shape.
• Register gradients by wrapping outputs in Tensor objects with _grad_fn pointing to backward classes like Conv2dBackward.
• Initialize convolutional weights using He (Kaiming) initialization with std = sqrt(2.0 / (in_channels * kernel_h * kernel_w)).
• Chain layers using Python's __call__ protocol to create readable model definitions that support automatic differentiation.

Frequently Asked Questions

Why use seven-nested loops instead of vectorized operations?

The explicit seven-nested loops in _convolve_loops prioritize educational clarity over computational performance. This implementation makes the convolution operation transparent, showing exactly how each output pixel results from the dot product of a kernel and a receptive field. Production frameworks use optimized BLAS libraries or im2col transformations for speed, but TinyTorch's approach ensures students understand the underlying mechanics before using abstractions.

How does the backward pass handle padding gradients?

The Conv2dBackward.apply method pads the gradient tensor when the forward pass used padding, then removes the padding before returning the final gradient with respect to the original input. This ensures that gradients flow correctly to the valid input regions while maintaining the mathematical correctness of the convolution operation's derivative.

What is the purpose of storing padded.shape in the backward function?

The Conv2dBackward class stores padded.shape (along with stride, padding, and kernel size) to reconstruct the exact geometry of the forward pass during backpropagation. These parameters allow the backward method to correctly place gradients in the padded tensor and then strip the padding to match the original input dimensions.

Can this implementation handle batch sizes greater than one?

Yes, the implementation explicitly handles batch processing. The _convolve_loops function iterates over the batch dimension, and the forward pass uses NumPy array slicing like out[:, oc, :, :] to apply the same convolutional filters across all samples in the batch simultaneously, maintaining the (batch, channel, height, width) layout throughout the network.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}