Let's put everything we have learned about PyTorch and CNNs into practice! In this project, you will build and train a deep learning classifier on the famous CIFAR-10 dataset.
The CIFAR-10 dataset contains 60,000 32x32 color images divided across 10 classes: airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. It is a classic dataset for testing new computer vision architectures.
Before feeding images into our neural network, we must normalize pixel values and convert them into PyTorch tensors. We chain these preprocessing steps using transforms.Compose:
from torchvision import transforms
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
Feeding all 60,000 images into a network at once is memory-prohibitive. We use a PyTorch DataLoaderto shuffle the dataset and serve images in small, manageable batches (e.g. batch size of 64).
For each batch, we perform a forward pass to calculate predictions, compute loss, zero existing gradients via optimizer.zero_grad(), perform a backward pass via loss.backward(), and update weights via optimizer.step().
Most bugs in a first CNN are shape mismatches, so work out the sizes on paper. A CIFAR-10 batch has shape (N, 3, 32, 32): N images, 3 color channels, 32 by 32 pixels. A 3x3 convolution with padding=1 keeps the height and width, and a 2x2 max-pool halves them.
import torch
import torch.nn as nn
import torch.nn.functional as F
class SmallCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1) # (N,16,32,32)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) # (N,32,16,16) after first pool
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32 * 8 * 8, 128) # 32 channels x 8 x 8 after second pool
self.fc2 = nn.Linear(128, 10) # one score per class
def forward(self, x):
x = self.pool(F.relu(self.conv1(x))) # (N,16,16,16)
x = self.pool(F.relu(self.conv2(x))) # (N,32,8,8)
x = torch.flatten(x, 1) # (N,2048)
x = F.relu(self.fc1(x)) # (N,128)
return self.fc2(x) # (N,10) raw scores (logits)
print(SmallCNN()(torch.zeros(4, 3, 32, 32)).shape) # torch.Size([4, 10])The final layer returns raw scores, called logits. Do not add a softmax: nn.CrossEntropyLoss applies it internally and expects logits plus integer class labels.
import torchvision
from torch.utils.data import DataLoader
from torchvision import transforms
import torch.optim as optim
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_set = torchvision.datasets.CIFAR10("./data", train=True, download=True, transform=transform)
test_set = torchvision.datasets.CIFAR10("./data", train=False, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
test_loader = DataLoader(test_set, batch_size=256)
model = SmallCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(2):
model.train()
for i, (images, labels) in enumerate(train_loader):
optimizer.zero_grad()
loss = criterion(model(images), labels)
loss.backward()
optimizer.step()
if i % 200 == 0:
print(epoch, i, loss.item())
model.eval()
correct = 0
with torch.no_grad():
for images, labels in test_loader:
correct += (model(images).argmax(dim=1) == labels).sum().item()
print("test accuracy:", correct / len(test_set))Normalize with mean 0.5 and standard deviation 0.5 maps pixel values from the range 0 to 1 into the range -1 to 1, which keeps inputs centered for training. The training set has 50,000 images and the test set 10,000, so batches of 64 give 782 steps per epoch. Random guessing scores 10%; a small network like this typically reaches well above that after two epochs, though the exact figure varies from run to run.
Track the training loss and the test accuracy after each epoch. If training loss keeps falling while test accuracy stalls or drops, the model is overfitting. Useful next steps, in rough order of effort: train for more epochs, add transforms.RandomHorizontalFlip() and RandomCrop(32, padding=4) to the training transform only, add dropout before fc1, or add a third convolution block. Make one change at a time so you know what helped.
optimizer.zero_grad(), so gradients accumulate across batches.fc1: recompute it whenever you change padding, pooling or channels.CrossEntropyLoss, which double-applies it.model.eval() and torch.no_grad(), which wastes memory and matters once you add dropout.Load and normalize the data, serve it in shuffled batches, define a CNN whose shapes you have checked by hand, train with the zero-grad, forward, backward, step cycle, and measure accuracy on held-out images.
To complete this project workshop, implement the following components inside the sandbox code panel: