Network Architecture#
In this notebook we build a compact, LeNet-inspired CNN for digit classification on the MNIST dataset. Along the way, we will track tensor shapes and connect convolution, pooling, and the classification head.
Image data#
Note
An image is a grid of pixels. A grayscale pixel contains one intensity value between 0 and 255; an RGB pixel contains red, green, and blue values between 0 and 255.
PyTorch represents an image as a 3D Tensor. The first dimension is the number of color channels, the second dimension is the height of the image, and the third dimension is the width of the image. Therefore, a grayscale image of size 28x28 pixels is represented as a tensor of shape (1, 28, 28), whereas a color image of the same size is represented as a tensor of shape (3, 28, 28). A batch adds a leading dimension, leading to a 4D tensor with logical shape
Image Batch = (batch, channels, height, width)
This memory layout is called channel-first or NCHW format. PyTorch also supports a channel-last or NHWC format, where the channel dimension comes after the height and width dimensions. The two formats are equivalent in terms of the logical order of dimensions, but they differ in how the data is laid out in memory. Throughout this course we keep the usual channel-first semantics.
MNIST dataset#
As usual, we load the MNIST dataset with a preprocessing pipeline that converts the images into PyTorch tensors and rescales them to have values between 0 and 1. We covered this topic in a previous tutorial, so we won’t go into too much detail here.
preprocess = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
train_ds = MNIST('.data', train=True, download=True, transform=preprocess)
test_ds = MNIST('.data', train=False, download=True, transform=preprocess)
100%|██████████| 9.91M/9.91M [00:00<00:00, 14.3MB/s]
100%|██████████| 28.9k/28.9k [00:00<00:00, 368kB/s]
100%|██████████| 1.65M/1.65M [00:00<00:00, 3.16MB/s]
100%|██████████| 4.54k/4.54k [00:00<00:00, 5.41MB/s]
Let’s take a look at the first image in the training set to confirm that the channel-first format is being used.
image, label = train_ds[0]
print('Image Size:', *image.shape)
Image Size: 1 28 28
The LeNet-5 architecture was proposed by Yann LeCun in 1998. It was one of the first convolutional neural networks and was designed to classify handwritten digits. The architecture consists of two convolutional layers alternating with max-pooling layers, then a flattening layer followed by three fully-connected layers. The input to the network is a grayscale image of size 28x28 pixels, whereas the output is a vector of size 10, representing the scores for each of the 10 classes. Let’s implement this architecture in PyTorch.
Creating a convolutional network#
LeNet-5, introduced by Yann LeCun and collaborators, was an influential early CNN for handwritten-character recognition. Our network follows its broad pattern, that is, two convolution/downsampling stages followed by a three-layer classifier, but it is not an exact reproduction. The original used 32×32 inputs, trainable subsampling, and different activations and connectivity. Here we use 28×28 MNIST images, ReLU, and max pooling to keep the implementation aligned with modern PyTorch practice.
Convolutional backbone#
A ConvNet generally starts off with convolutional and pooling layers. In our case, we stack two convolutional layers, alternated with pooling layers.
Convolutional layers require us to specify the number of input channels and the number of output channels. The latter corresponds to the number of trainable filters that will be convolved with the input. We also need to specify the size of each filter, which is 5x5 in LeNet-5. We use a padding of 2 in the first layer, so that the spatial dimensions of the input tensor will remain the same after the convolution.
Pooling layers require us to specify the spatial dimensions of the pooling window, and optionally the strides, which defaults to the same value as the pooling window. In this case, we use a 2x2 window, which means that the spatial dimensions of the input tensor will be halved after the pooling operation.
class Backbone(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 6, kernel_size=5, padding=2)
self.conv2 = torch.nn.Conv2d(6, 16, kernel_size=5)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
return x
Let’s display the output of this convolutional backbone on a batch of MNIST images.
Output Size: 7 16 5 5
The backbone returns a 4D tensor of shape (batch, 16, 5, 5). Channels increase from 1 to 16 while spatial dimensions shrink, a common CNN pattern that trades spatial detail for a richer set of learned features.
For one spatial dimension, a convolution or pooling layer produces the following output size, given the input size, kernel size, padding, and stride.
output = floor((input + 2 × padding - dilation × (kernel - 1) - 1) / stride + 1)
Applying the formula at each stage gives the shapes reported in the table below. Here dilation is 1.
Layer |
Output Shape |
Calculation |
|---|---|---|
Input |
|
|
Conv1 |
|
|
Pool1 |
|
|
Conv2 |
|
|
Pool2 |
|
|
Note
Padding affects what happens near image boundaries. Border activations partly depend on padded values rather than observed pixels.
Classification head#
A ConvNet designed for image classification usually ends with a few fully-connected layers. These layers process the output of the convolutional backbone and generate the final class predictions. However, the output of the convolutional backbone is a 3D tensor, whereas fully-connected layers expect 1D vectors as input. A flattening operation is therefore required between the convolutional backbone and the fully-connected layers.
Note
Modern ConvNets often use adaptive global pooling to reduce the spatial dimensions of the feature maps to 1x1. This allows the network to handle inputs of varying sizes and reduces the number of parameters in the fully-connected layers.
The final layer returns ten logits, one per digit. We deliberately omit softmax because CrossEntropyLoss applies the required log-softmax operation internally in a numerically stable way.
class ClassificationHead(torch.nn.Module):
def __init__(self):
super().__init__()
self.fc1 = torch.nn.Linear(16*5*5, 120)
self.fc2 = torch.nn.Linear(120, 84)
self.fc3 = torch.nn.Linear(84, 10)
def forward(self, x):
x = torch.flatten(x, start_dim=1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
Final network#
The complete network composes the convolutional backbone with the classification head. Keeping the two parts separate makes their roles, and the shape transition between them, easy to inspect.
class LeNet5(torch.nn.Module):
def __init__(self):
super().__init__()
self.backbone = Backbone()
self.head = ClassificationHead()
def forward(self, x):
x = self.backbone(x)
x = self.head(x)
return x
This is the resulting module hierarchy.
LeNet5(
(backbone): Backbone(
(conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
)
(head): ClassificationHead(
(fc1): Linear(in_features=400, out_features=120, bias=True)
(fc2): Linear(in_features=120, out_features=84, bias=True)
(fc3): Linear(in_features=84, out_features=10, bias=True)
)
)
As a sanity check, pass random values with the same shape as a batch of MNIST images through the model.
Output Size: 7 10
Training LeNet-5#
Now, let’s train the CNN on the MNIST digits using the Trainer class from the training.py file.
from training import Trainer
Training loop#
We will train the network for 5 epochs using the cross-entropy loss function.
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_ds, batch_size=512)
model = LeNet5()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()
epochs = 5
trainer = Trainer()
history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs)
===== Training on cpu device =====
Epoch 1/5: 100%|██████████| 938/938 [00:24<00:00, 38.83it/s, train_loss=0.2947]
Epoch 2/5: 100%|██████████| 938/938 [00:28<00:00, 32.47it/s, train_loss=0.0797]
Epoch 3/5: 100%|██████████| 938/938 [00:29<00:00, 32.05it/s, train_loss=0.0565]
Epoch 4/5: 100%|██████████| 938/938 [00:30<00:00, 31.20it/s, train_loss=0.0437]
Epoch 5/5: 100%|██████████| 938/938 [00:29<00:00, 31.93it/s, train_loss=0.0371]
Evaluation#
Let’s also evaluate the model on the test data.
from torcheval.metrics import MulticlassAccuracy
trainer.set_metrics(accuracy=MulticlassAccuracy())
ans = trainer.eval(model, test_loader)
print(f'Test Accuracy: {ans["accuracy"]:.2%}')
Test Accuracy: 98.61%
Summary#
In this tutorial, we built a convolutional neural network for MNIST digit classification. Inspired by LeNet-5 architecture, the network consists of two convolutional layers with max pooling, followed by a flatten layer and three fully-connected layers. The network takes (batch, 1, 28, 28) images as input and produces ten output logits.
The comparison with the earlier MLP is useful, but a single run is not enough to claim a universal accuracy advantage. A careful comparison keeps the split and training budget fixed and reports results across several seeds. The architectural advantage is clearer: unlike an MLP on flattened pixels, the CNN explicitly preserves and exploits spatial structure.