PyTorch Convolutional Neural Networks

Real-World Scenario

A manufacturing company runs visual quality inspection on an assembly line. A CNN trained on 50,000 defect images classifies parts as pass/fail at 120ms per image — fast enough for real-time production line decisions. CNNs are the workhorse behind every image recognition system in production.

How Convolutions Work

import torch
import torch.nn as nn

# A 2D convolution layer
# in_channels=1 (grayscale), out_channels=32 (learn 32 filters), kernel_size=3
conv = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)

# Input: batch of 8 grayscale 28x28 images
x = torch.randn(8, 1, 28, 28)
output = conv(x)
print(f"Input:  {x.shape}")    # (8, 1, 28, 28)
print(f"Output: {output.shape}")  # (8, 32, 28, 28) — padding=1 preserves spatial dims

# Pooling — reduces spatial dimensions, extracts dominant features
pool = nn.MaxPool2d(kernel_size=2, stride=2)
pooled = pool(output)
print(f"After pooling: {pooled.shape}")  # (8, 32, 14, 14)

Building a CNN for MNIST

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

class MNISTConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        # Conv block 1: 1 → 32 feature maps, 28x28 → 14x14
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.bn1   = nn.BatchNorm2d(32)

        # Conv block 2: 32 → 64 feature maps, 14x14 → 7x7
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2   = nn.BatchNorm2d(64)

        # Classifier head
        self.dropout = nn.Dropout(0.5)
        self.fc1     = nn.Linear(64 * 7 * 7, 256)
        self.fc2     = nn.Linear(256, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Block 1
        x = F.relu(self.bn1(self.conv1(x)))  # (B, 32, 28, 28)
        x = F.max_pool2d(x, 2)               # (B, 32, 14, 14)

        # Block 2
        x = F.relu(self.bn2(self.conv2(x)))  # (B, 64, 14, 14)
        x = F.max_pool2d(x, 2)               # (B, 64, 7, 7)

        # Flatten and classify
        x = x.view(x.size(0), -1)            # (B, 64*7*7 = 3136)
        x = F.relu(self.fc1(self.dropout(x)))
        return self.fc2(x)                   # logits: (B, 10)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = MNISTConvNet().to(device)

# Count parameters
params = sum(p.numel() for p in model.parameters())
print(f"Parameters: {params:,}")  # ~820k

# Data loading with augmentation
transform_train = transforms.Compose([
    transforms.RandomRotation(10),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])

train_ds = datasets.MNIST("./data", train=True,  download=True, transform=transform_train)
test_ds  = datasets.MNIST("./data", train=False, download=True, transform=transform_test)
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True,  num_workers=2, pin_memory=True)
test_loader  = DataLoader(test_ds,  batch_size=256, shuffle=False, num_workers=2, pin_memory=True)

optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
scheduler = optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=1e-2, steps_per_epoch=len(train_loader), epochs=10
)
criterion = nn.CrossEntropyLoss()

def epoch_pass(model, loader, optimizer=None, criterion=None):
    training = optimizer is not None
    model.train() if training else model.eval()
    total_loss, correct, total = 0.0, 0, 0

    with torch.set_grad_enabled(training):
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            logits = model(images)
            if training:
                loss = criterion(logits, labels)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                scheduler.step()
                total_loss += loss.item() * len(labels)
            correct += (logits.argmax(1) == labels).sum().item()
            total   += len(labels)

    return total_loss / total if training else None, correct / total


for epoch in range(10):
    _, train_acc = epoch_pass(model, train_loader, optimizer, criterion)
    _, test_acc  = epoch_pass(model, test_loader)
    print(f"Epoch {epoch+1:2d}: Train={train_acc:.3f}  Test={test_acc:.3f}")

Transfer Learning with Pretrained Models

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load ResNet18 pretrained on ImageNet — 11M parameters already trained
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

# Freeze all layers — we only want to train the final classifier
for param in model.parameters():
    param.requires_grad = False

# Replace the final layer for our task (e.g., 5-class defect classification)
num_classes = 5
model.fc = nn.Linear(model.fc.in_features, num_classes)  # new layer has requires_grad=True

model = model.to(device)

# Only optimize the new final layer — much faster, needs less data
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

# ImageNet normalization — required for pretrained ResNet
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# After training the head, unfreeze and fine-tune the whole network
def fine_tune_all(model, new_lr=1e-5):
    """Unfreeze all layers for end-to-end fine-tuning."""
    for param in model.parameters():
        param.requires_grad = True
    return torch.optim.Adam(model.parameters(), lr=new_lr)

# Verify only the new layer has gradients initially
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:,} parameters ({100*trainable/total:.1f}%)")