PyTorch Convolutional Neural Networks

Real-World Scenario

A manufacturing company runs visual quality inspection on an assembly line: a CNN trained on photos of defective and good parts classifies each new part as pass or fail. CNNs are the workhorse behind image recognition in production, and the same building blocks — convolutions, pooling, a small classifier head — show up whether the images are digits or circuit boards.

Convolution Output Shapes

A Conv2d layer takes (batch, channels, height, width) and produces out_channels feature maps. Padding and stride decide their size.

import torch
import torch.nn as nn

x = torch.randn(8, 1, 28, 28)           # (batch, channels, height, width)

same = nn.Conv2d(1, 32, kernel_size=3, padding=1)
valid = nn.Conv2d(1, 32, kernel_size=3, padding=0)
strided = nn.Conv2d(1, 32, kernel_size=3, padding=1, stride=2)
pool = nn.MaxPool2d(kernel_size=2)

print("padding=1:", same(x).shape)
print("padding=0:", valid(x).shape)
print("stride=2: ", strided(x).shape)
print("pooled:   ", pool(same(x)).shape)

$ python conv_shapes.py
padding=1: torch.Size([8, 32, 28, 28])
padding=0: torch.Size([8, 32, 26, 26])
stride=2:  torch.Size([8, 32, 14, 14])
pooled:    torch.Size([8, 32, 14, 14])

With a 3×3 kernel, padding=1 adds a one-pixel border so the output stays 28×28; padding=0 loses one pixel on each side, giving 26×26. Both a stride of 2 and a 2×2 max-pool halve the resolution to 14×14 — two different ways to downsample. In every case the channel dimension becomes 32, one map per learned filter.

Weight Sharing

import torch.nn as nn

conv = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
dense = nn.Linear(28 * 28, 32 * 26 * 26)   # same input/output sizes, no sharing

print(conv.weight.shape)
print(f"conv:  {sum(p.numel() for p in conv.parameters()):>12,}")
print(f"dense: {sum(p.numel() for p in dense.parameters()):>12,}")

$ python conv_params.py
torch.Size([32, 1, 3, 3])
conv:           320
dense:   16,981,120

The conv layer’s weight is (32, 1, 3, 3): 32 filters of 3×3 over one input channel, plus 32 biases — 320 numbers. A fully connected layer producing the same 32×26×26 output needs almost 17 million, because it learns a separate weight for every input-output pair. The convolution reuses each small filter at every position, which is both why it is cheap and why a pattern it learns to detect in one corner is detected everywhere.

What a Filter Computes

A filter is a small grid of weights; at each position the output is the sum of the weights times the pixels underneath. Set one by hand to see it:

import torch
import torch.nn as nn

image = torch.zeros(1, 1, 5, 6)
image[..., 3:] = 1.0                     # dark left half, bright right half
print(image[0, 0])

conv = nn.Conv2d(1, 1, kernel_size=3, bias=False)
with torch.no_grad():
    conv.weight[:] = torch.tensor([[-1., 0., 1.],
                                   [-1., 0., 1.],
                                   [-1., 0., 1.]])   # right minus left
    print(conv(image)[0, 0])

$ python edge_filter.py
tensor([[0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1.]])
tensor([[0., 3., 3., 0.],
        [0., 3., 3., 0.],
        [0., 3., 3., 0.]])

The kernel subtracts the left column from the right column of each 3×3 window. Over flat regions — all dark or all bright — that is zero. Where a window straddles the dark-to-bright boundary, three bright pixels on the right and three dark on the left give 3. The output is 3×4 because no padding was used. A trained CNN learns kernels like this from data instead of having them written by hand.

Building a CNN for MNIST

Two convolution blocks shrink 28×28 digits to 64 maps of 7×7; a small classifier head turns those into 10 logits. Saved as mnist_model.py, the model is imported by the training script below.

import torch
import torch.nn as nn

class MNISTConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2),
        )
        self.head = nn.Sequential(
            nn.Flatten(), nn.Dropout(0.5), nn.Linear(64 * 7 * 7, 256), nn.ReLU(),
            nn.Linear(256, 10),
        )

    def forward(self, x):
        return self.head(self.features(x))

if __name__ == "__main__":
    model = MNISTConvNet()
    x = torch.randn(4, 1, 28, 28)
    for layer in model.features:
        x = layer(x)
        if isinstance(layer, (nn.Conv2d, nn.MaxPool2d)):
            print(f"{type(layer).__name__:10s} -> {tuple(x.shape)}")
    print(f"parameters: {sum(p.numel() for p in model.parameters()):,}")

$ python mnist_model.py
Conv2d     -> (4, 32, 28, 28)
MaxPool2d  -> (4, 32, 14, 14)
Conv2d     -> (4, 64, 14, 14)
MaxPool2d  -> (4, 64, 7, 7)
parameters: 824,650

Each MaxPool2d halves height and width, and each convolution changes the channel count. That trace is where the 64 * 7 * 7 in the first Linear layer comes from — if you change the input size or add a pooling layer, that number must change too, or the forward pass fails with a shape mismatch. Most of the 824,650 parameters sit in that one Linear(3136, 256) layer, not in the convolutions.

Training on MNIST

import time
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from mnist_model import MNISTConvNet

torch.manual_seed(0)
tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train = DataLoader(datasets.MNIST("data", train=True, download=True, transform=tf),
                   batch_size=128, shuffle=True, num_workers=2)
test = DataLoader(datasets.MNIST("data", train=False, transform=tf), batch_size=512)

model = MNISTConvNet()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
criterion = nn.CrossEntropyLoss()

start = time.perf_counter()
for epoch in range(1, 4):
    model.train()
    for images, labels in train:
        optimizer.zero_grad()
        criterion(model(images), labels).backward()
        optimizer.step()
    model.eval()
    with torch.no_grad():
        correct = sum((model(x).argmax(1) == y).sum().item() for x, y in test)
    print(f"epoch {epoch}: test accuracy {correct / len(test.dataset):.4f}")
print(f"training time: {time.perf_counter() - start:.0f} s")

$ python mnist_train.py
epoch 1: test accuracy 0.9854
epoch 2: test accuracy 0.9903
epoch 3: test accuracy 0.9900
training time: 130 s

The first run downloads MNIST (about 12 MB) into data/; later runs reuse it. Three epochs on a 4-core CPU reach 99% test accuracy in the training time shown. Normalize((0.1307,), (0.3081,)) rescales pixels using MNIST’s own mean and standard deviation, and num_workers=2 decodes batches in background processes while the model trains. Accuracy levels off after the second epoch — epoch 3 is a hair below epoch 2 — so squeezing out more takes a learning-rate schedule and augmentation such as transforms.RandomRotation(10), not just more epochs at a fixed rate.

Transfer Learning with a Pretrained ResNet

Rather than training from scratch, take a network already trained on ImageNet and replace only its last layer.

import torch.nn as nn
from torchvision import models

# weights=None builds the architecture only; use ResNet18_Weights.IMAGENET1K_V1 for real work
model = models.resnet18(weights=None)
print(model.fc)

for param in model.parameters():
    param.requires_grad = False                  # freeze the backbone

model.fc = nn.Linear(model.fc.in_features, 5)   # new head, trainable by default

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"trainable: {trainable:,} / {total:,} ({100 * trainable / total:.2f}%)")
print(models.ResNet18_Weights.IMAGENET1K_V1.transforms())

$ python transfer.py
Linear(in_features=512, out_features=1000, bias=True)
trainable: 2,565 / 11,179,077 (0.02%)
ImageClassification(
    crop_size=[224]
    resize_size=[256]
    mean=[0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]
    interpolation=InterpolationMode.BILINEAR
)

ResNet-18’s final layer maps 512 features to ImageNet’s 1000 classes. After freezing every existing parameter and swapping in a new 5-class head, only 512·5 + 5 = 2,565 of 11.2 million parameters train — which is why a frozen backbone learns quickly from a few hundred images. The parameter counts do not depend on the weights, so this example uses weights=None and skips the 45 MB download; in real use pass weights=models.ResNet18_Weights.IMAGENET1K_V1. The weights object also carries the exact preprocessing the network was trained with: resize to 256, crop to 224, and normalize with ImageNet’s channel means and standard deviations.

Fine-Tuning with Parameter Groups

Once the head has trained, you can unfreeze everything and fine-tune — but the pretrained layers need much smaller steps than the new head.

import torch
import torch.nn as nn
from torchvision import models

model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, 5)

head = list(model.fc.parameters())
head_ids = {id(p) for p in head}
backbone = [p for p in model.parameters() if id(p) not in head_ids]

optimizer = torch.optim.AdamW([
    {"params": backbone, "lr": 1e-5},     # pretrained layers: small steps
    {"params": head, "lr": 1e-3},         # new layer: larger steps
])
for group in optimizer.param_groups:
    print(len(group["params"]), "tensors, lr", group["lr"])

$ python fine_tune.py
60 tensors, lr 1e-05
2 tensors, lr 0.001

An optimizer accepts a list of parameter groups, each with its own settings. The 60 backbone tensors get a learning rate of 1e-5, small enough not to destroy what they learned on ImageNet, while the head’s weight and bias train at 1e-3. Identifying parameters by id() avoids comparing tensors with ==, which would compare their values element by element.

Practice

import torch
import torch.nn as nn

conv = nn.Conv2d(3, 16, kernel_size=5, stride=2, padding=2)
x = torch.randn(1, 3, 32, 32)

predicted = (32 + 2 * 2 - 5) // 2 + 1        # (size + 2*padding - kernel) // stride + 1
print(predicted)
print(conv(x).shape)

$ python p1_formula.py
16
torch.Size([1, 16, 16, 16])

import torch
import torch.nn as nn

image = torch.zeros(1, 1, 6, 5)
image[..., 3:, :] = 1.0                  # dark top, bright bottom

conv = nn.Conv2d(1, 1, kernel_size=3, bias=False)
with torch.no_grad():
    conv.weight[:] = torch.tensor([[-1., -1., -1.],
                                   [ 0.,  0.,  0.],
                                   [ 1.,  1.,  1.]])   # below minus above
    print(conv(image)[0, 0])

$ python p2_horizontal.py
tensor([[0., 0., 0.],
        [3., 3., 3.],
        [3., 3., 3.],
        [0., 0., 0.]])

import torch
import torch.nn as nn
from torchvision import models

model = models.resnet18(weights=None)
model.fc = nn.Identity()                  # drop the classifier, keep the features
model.eval()

with torch.no_grad():
    features = model(torch.randn(4, 3, 224, 224))
print(features.shape)

$ python p3_embeddings.py
torch.Size([4, 512])

import torch
import torch.nn as nn
from mnist_model import MNISTConvNet

model = MNISTConvNet()
model.head = nn.Sequential(
    nn.AdaptiveAvgPool2d(1),              # average each 7x7 map down to one number
    nn.Flatten(),
    nn.Linear(64, 10),
)
print(model(torch.randn(2, 1, 28, 28)).shape)
print(f"parameters: {sum(p.numel() for p in model.parameters()):,}")

$ python p4_gap.py
torch.Size([2, 10])
parameters: 19,658