PyTorch CNN for Images

This page builds a complete image classification CNN in PyTorch from data loading to evaluation. The example classifies handwritten digits from the MNIST dataset — ten classes (0 through 9), 28×28 grayscale images. Every step follows the same pattern used in production image models, just at a smaller scale.

The Full Pipeline

Raw Images on Disk
    │
    ▼ Dataset + DataLoader
Batches of Tensors
    │
    ▼ CNN (conv → pool → conv → pool → flatten → FC)
Class Scores
    │
    ▼ CrossEntropyLoss
Loss Value
    │
    ▼ backward() + optimizer.step()
Updated Weights
    │
    ▼ Evaluate on validation set
Accuracy

Step 1: Load and Prepare Data

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

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))   # MNIST mean and std
])

train_data = datasets.MNIST('./data', train=True,  download=True, transform=transform)
test_data  = datasets.MNIST('./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_data,  batch_size=64, shuffle=False)

Step 2: Define the CNN Architecture

class DigitCNN(nn.Module):
    def __init__(self):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),   # [B,32,28,28]
            nn.ReLU(),
            nn.MaxPool2d(2, 2),                           # [B,32,14,14]

            nn.Conv2d(32, 64, kernel_size=3, padding=1),  # [B,64,14,14]
            nn.ReLU(),
            nn.MaxPool2d(2, 2)                            # [B,64, 7, 7]
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 10)
        )

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = DigitCNN().to(device)
print(model)

Architecture Shape Flow

Input:              [64, 1, 28, 28]   ← batch=64, 1 channel, 28×28

Conv2d(1→32,3,p=1): [64, 32, 28, 28]  ← 32 feature maps
ReLU
MaxPool2d(2):       [64, 32, 14, 14]  ← halved

Conv2d(32→64,3,p=1):[64, 64, 14, 14]  ← 64 feature maps
ReLU
MaxPool2d(2):       [64, 64, 7, 7]    ← halved again

Flatten:            [64, 3136]         ← 64×7×7 = 3136
Linear(3136→128):   [64, 128]
ReLU + Dropout
Linear(128→10):     [64, 10]          ← 10 class scores

Step 3: Training Loop

optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn   = nn.CrossEntropyLoss()

def train_epoch(model, loader, optimizer, loss_fn, device):
    model.train()
    total_loss, correct = 0, 0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss    = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        correct    += (outputs.argmax(1) == labels).sum().item()

    avg_loss = total_loss / len(loader)
    accuracy = correct / len(loader.dataset)
    return avg_loss, accuracy

Step 4: Evaluation Loop

def evaluate(model, loader, loss_fn, device):
    model.eval()
    total_loss, correct = 0, 0

    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss    = loss_fn(outputs, labels)
            total_loss += loss.item()
            correct    += (outputs.argmax(1) == labels).sum().item()

    return total_loss / len(loader), correct / len(loader.dataset)

Step 5: Run Training

for epoch in range(10):
    train_loss, train_acc = train_epoch(model, train_loader, optimizer, loss_fn, device)
    val_loss,   val_acc   = evaluate(model, test_loader, loss_fn, device)

    print(f"Epoch {epoch+1:2d}  "
          f"Train Loss: {train_loss:.4f}  Train Acc: {train_acc:.4f}  "
          f"Val Loss: {val_loss:.4f}  Val Acc: {val_acc:.4f}")

Expected output after 10 epochs:

Epoch  1  Train Loss: 0.2341  Train Acc: 0.9289  Val Loss: 0.0712  Val Acc: 0.9776
Epoch  5  Train Loss: 0.0631  Train Acc: 0.9808  Val Loss: 0.0443  Val Acc: 0.9865
Epoch 10  Train Loss: 0.0421  Train Acc: 0.9870  Val Loss: 0.0381  Val Acc: 0.9890

Step 6: Make Predictions

model.eval()
images, labels = next(iter(test_loader))
images = images.to(device)

with torch.no_grad():
    outputs = model(images)
    predicted = outputs.argmax(dim=1)

print("True labels:     ", labels[:8].tolist())
print("Predicted labels:", predicted[:8].cpu().tolist())

Saving and Loading the Trained Model

# Save
torch.save(model.state_dict(), 'digit_cnn.pth')

# Load
model_loaded = DigitCNN().to(device)
model_loaded.load_state_dict(torch.load('digit_cnn.pth'))
model_loaded.eval()

Summary

Building a CNN for image classification follows six steps: load data with transforms and DataLoader, define the CNN architecture using Conv2d and pooling blocks, set up a loss function and optimizer, write a training loop that iterates over batches, write an evaluation loop with no_grad, and run training for several epochs while monitoring both training and validation metrics. This same pipeline scales directly from MNIST to CIFAR-10, ImageNet, and custom datasets — only the architecture size and data loading code change.