PyTorch CNN: Convolutional Neural Networks

A CNN in PyTorch (convolutional neural network) is the classic architecture for images. Instead of connecting every pixel to every neuron, a convolution slides small learnable filters across the image, detecting edges and textures that are useful anywhere in the picture. This uses far fewer parameters and respects the 2-D structure of images.

Conv2d and Output Shapes

nn.Conv2d(in_channels, out_channels, kernel_size) expects input of shape [batch, channels, height, width]. The output size per side is (input - kernel + 2 x padding) / stride + 1.

import torch
from torch import nn

x = torch.randn(8, 1, 28, 28)
conv = nn.Conv2d(1, 32, kernel_size=3, padding=1)
print(conv(x).shape)
print(nn.Conv2d(1, 32, kernel_size=3)(x).shape)

torch.Size([8, 32, 28, 28])
torch.Size([8, 32, 26, 26])

padding=1 with a 3x3 kernel keeps the size the same; without padding each side shrinks by 2.

Pooling

pool = nn.MaxPool2d(2)
print(pool(conv(x)).shape)

torch.Size([8, 32, 14, 14])

Max pooling keeps the largest value in each 2x2 window, halving height and width. This reduces computation and makes features slightly more robust to small shifts.

A Small CNN for MNIST

class DigitCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # -> 32x14x14
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # -> 64x7x7
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, 10),
        )

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

model = DigitCNN()
print(model(torch.randn(2, 1, 28, 28)).shape)
print(sum(p.numel() for p in model.parameters()))

torch.Size([2, 10])
421642

Train it with exactly the same loop as in the MNIST lesson. After a few epochs it typically reaches around 99% test accuracy.

Training Output

# using train_dl, test_dl, evaluate() from the MNIST lesson
for epoch in range(3):
    ...  # same training loop
    print(f"epoch {epoch + 1}: test acc {evaluate():.4f}")

epoch 1: test acc 0.9842
epoch 2: test acc 0.9881
epoch 3: test acc 0.9907
(sample output)

CNN Building Blocks to Know

• nn.BatchNorm2d normalizes activations and makes deeper CNNs train faster.
• nn.AdaptiveAvgPool2d(1) squeezes any spatial size to 1x1, so the network accepts different image sizes.
• Residual (skip) connections, as in ResNet, allow very deep networks.
• Data augmentation (flips, crops, color jitter) is crucial for real photos.

Frequently asked questions

How do I calculate the output size of Conv2d?

Use floor((input + 2 x padding - kernel_size) / stride) + 1 for each spatial dimension.

Why use a CNN instead of a fully connected network for images?

CNNs share filter weights across the image, so they need far fewer parameters, learn local patterns and generalize better to shifted or slightly changed images.

What does in_channels mean in Conv2d?

The number of channels in the input: 1 for grayscale, 3 for RGB images, or the out_channels of the previous convolution.