Eric Campbell 5 min read Machine Learning

How to Implement a CNN in PyTorch

Convolutional Neural Networks revolutionized computer vision by automatically learning hierarchical feature representations from raw pixel data. Unlike traditional neural networks that treat images...

Key Insights

  • PyTorch’s dynamic computation graph and intuitive API make it ideal for prototyping CNNs, with the nn.Module class providing a clean abstraction for defining network architectures
  • A production-ready CNN implementation requires careful attention to data preprocessing, batch normalization, and regularization techniques like dropout to prevent overfitting
  • The training loop follows a consistent pattern: forward pass, loss computation, backpropagation, and optimizer step—understanding this cycle is essential for debugging and optimization

Introduction to CNNs and PyTorch

Convolutional Neural Networks revolutionized computer vision by automatically learning hierarchical feature representations from raw pixel data. Unlike traditional neural networks that treat images as flat vectors, CNNs preserve spatial relationships through specialized layers that apply learnable filters across the input.

PyTorch excels at CNN implementation due to its imperative programming model and transparent tensor operations. You define networks as Python classes, making the architecture explicit and debuggable. The framework handles automatic differentiation through its autograd system, letting you focus on model design rather than gradient calculations.

A typical CNN consists of three layer types: convolutional layers extract features through learned filters, pooling layers reduce spatial dimensions while retaining important information, and fully connected layers perform final classification. Modern architectures stack these components with careful attention to receptive fields and feature map dimensions.

Setting Up the Environment

Install PyTorch with CUDA support if you have a GPU available. The official PyTorch website provides platform-specific installation commands.

pip install torch torchvision torchaudio

Here are the essential imports for a complete CNN implementation:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
import numpy as np
from tqdm import tqdm

We’ll use CIFAR-10, a dataset of 60,000 32x32 color images across 10 classes. It’s complex enough to demonstrate CNN capabilities without requiring excessive compute resources:

# Define transforms
transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# Load datasets
train_dataset = CIFAR10(root='./data', train=True, download=True, 
                        transform=transform_train)
test_dataset = CIFAR10(root='./data', train=False, download=True, 
                       transform=transform_test)

Defining the CNN Architecture

Subclass nn.Module to define your network. The __init__ method declares layers, and forward() specifies how data flows through them. This separation allows PyTorch to build the computation graph automatically.

class CNN(nn.Module):
    def __init__(self, num_classes=10):
        super(CNN, self).__init__()
        
        # First convolutional block
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, 
                               kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Second convolutional block
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.bn4 = nn.BatchNorm2d(128)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Third convolutional block
        self.conv5 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.bn5 = nn.BatchNorm2d(256)
        self.conv6 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.bn6 = nn.BatchNorm2d(256)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Fully connected layers
        self.fc1 = nn.Linear(256 * 4 * 4, 512)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, num_classes)
        
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        # Block 1
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.pool1(x)
        
        # Block 2
        x = self.relu(self.bn3(self.conv3(x)))
        x = self.relu(self.bn4(self.conv4(x)))
        x = self.pool2(x)
        
        # Block 3
        x = self.relu(self.bn5(self.conv5(x)))
        x = self.relu(self.bn6(self.conv6(x)))
        x = self.pool3(x)
        
        # Flatten and classify
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        
        return x

Batch normalization layers stabilize training by normalizing activations. Padding preserves spatial dimensions through convolutions. The dropout layer prevents overfitting by randomly zeroing activations during training.

Preparing the Data Pipeline

DataLoaders handle batching, shuffling, and parallel data loading. Proper configuration significantly impacts training efficiency:

batch_size = 128
num_workers = 4  # Adjust based on CPU cores

train_loader = DataLoader(train_dataset, batch_size=batch_size, 
                         shuffle=True, num_workers=num_workers, 
                         pin_memory=True)

test_loader = DataLoader(test_dataset, batch_size=batch_size, 
                        shuffle=False, num_workers=num_workers, 
                        pin_memory=True)

Set pin_memory=True when using CUDA to speed up host-to-device transfers. Use num_workers > 0 to parallelize data loading, preventing the GPU from idling while waiting for batches.

Training the Model

The training loop implements the core learning algorithm. Move your model and data to the GPU if available:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CNN(num_classes=10).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', 
                                                  factor=0.5, patience=5)

def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, targets in tqdm(loader, desc='Training'):
        inputs, targets = inputs.to(device), targets.to(device)
        
        # Forward pass
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        # Backward pass
        loss.backward()
        optimizer.step()
        
        # Statistics
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    
    epoch_loss = running_loss / len(loader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

def validate(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, targets in tqdm(loader, desc='Validation'):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    
    epoch_loss = running_loss / len(loader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

# Training loop
num_epochs = 50
best_acc = 0

for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch(model, train_loader, criterion, 
                                       optimizer, device)
    val_loss, val_acc = validate(model, test_loader, criterion, device)
    
    scheduler.step(val_loss)
    
    print(f'Epoch {epoch+1}/{num_epochs}')
    print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')
    print(f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')
    
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), 'best_model.pth')

Always call optimizer.zero_grad() before backpropagation—PyTorch accumulates gradients by default. Use model.train() and model.eval() to toggle dropout and batch normalization behavior.

Evaluation and Inference

After training, evaluate on the test set and examine per-class performance:

def evaluate_model(model, loader, device, class_names):
    model.eval()
    all_preds = []
    all_targets = []
    
    with torch.no_grad():
        for inputs, targets in loader:
            inputs = inputs.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            
            all_preds.extend(predicted.cpu().numpy())
            all_targets.extend(targets.numpy())
    
    accuracy = 100. * np.sum(np.array(all_preds) == np.array(all_targets)) / len(all_targets)
    print(f'Test Accuracy: {accuracy:.2f}%')
    
    return all_preds, all_targets

# Load best model
model.load_state_dict(torch.load('best_model.pth'))
class_names = train_dataset.classes

predictions, targets = evaluate_model(model, test_loader, device, class_names)

# Single image prediction
def predict_image(model, image_tensor, device, class_names):
    model.eval()
    with torch.no_grad():
        image_tensor = image_tensor.unsqueeze(0).to(device)
        output = model(image_tensor)
        probabilities = torch.softmax(output, dim=1)
        confidence, predicted = probabilities.max(1)
        
    return class_names[predicted.item()], confidence.item()

Tips for Optimization

Learning rate scheduling prevents the optimizer from overshooting minima. The ReduceLROnPlateau scheduler decreases the learning rate when validation loss plateaus:

scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

Implement gradient clipping for training stability:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Add dropout to convolutional layers for stronger regularization:

self.dropout_conv = nn.Dropout2d(0.2)

Monitor GPU memory usage and reduce batch size if you encounter out-of-memory errors. Use mixed precision training with torch.cuda.amp for faster training on modern GPUs:

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Always validate your model on a held-out test set. Cross-validation provides more robust performance estimates for smaller datasets. Track experiments with tools like TensorBoard or Weights & Biases to compare hyperparameter configurations systematically.

Liked this? There's more.

Every week: one practical technique, explained simply, with code you can use immediately.