Image Classification with CNNs in PyTorch

How to build, train, and evaluate a Convolutional Neural Network (CNN) for image classification tasks in Python using PyTorch and torchvision.

Computer vision has unlocked unprecedented capabilities in artificial intelligence, ranging from autonomous vehicle navigation and medical imaging diagnostics to facial recognition and automated quality control. At the core of computer vision breakthroughs lies the **Convolutional Neural Network (CNN)**, a specialized class of deep neural network explicitly designed to process grid-like pixel data such as images.

While standard fully connected feedforward networks struggle with high-dimensional image data due to parameter explosion and loss of spatial context, CNNs leverage localized receptive fields, weight sharing, and spatial pooling to extract hierarchical visual features automatically. This comprehensive guide walks you through building, training, and evaluating a complete image classification CNN from scratch using **PyTorch** and **torchvision**.

1. Understanding Convolutional Neural Network Architecture

A typical image classification CNN is composed of a sequence of specialized architectural layers designed to transform raw pixel arrays into high-level class probabilities.

Convolutional Layers (`nn.Conv2d`): The core building block of a CNN. These layers apply a set of learnable filters (kernels) that slide across the input image, performing element-wise multiplication and summation to extract local spatial features such as edges, corners, and textures.
Activation Functions (`ReLU`): Applied element-wise after convolutions, Rectified Linear Units introduce non-linearity into the network, allowing the model to learn complex, non-linear decision boundaries.
Pooling Layers (`nn.MaxPool2d`): Spatial pooling (typically max pooling) downsamples feature maps by taking the maximum value over local windows. This reduces spatial dimensions, controls overfitting, and provides translation invariance.
Fully Connected Layers (`nn.Linear`): Positioned at the tail end of the network, fully connected layers map the flattened high-level feature representations into final class logits for classification.

2. Environment Setup and Dataset Preparation

Let's set up our Python development environment and load the benchmark dataset. For this tutorial, we will use the classic **CIFAR-10** dataset, consisting of 60,000 32x32 color images across 10 distinct classes.

              
C

            # Create and activate a virtual environment
python -m venv venv
source venv/bin/activate  # On Windows use: venv\Scripts\activate

# Install PyTorch, torchvision, and matplotlib
pip install torch torchvision matplotlib

Next, let's write our data loading and preprocessing pipeline using `torchvision.datasets` and `DataLoader`:

              
            
C

            
            import torch
import torchvision
import torchvision.transforms as transforms

# Define data augmentation and normalization transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# Download and load training dataset
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True,
    download=True, transform=transform
)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=64,
    shuffle=True, num_workers=2
)

# Download and load test dataset
testset = torchvision.datasets.CIFAR10(
    root='./data', train=False,
    download=True, transform=transform
)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=64,
    shuffle=False, num_workers=2
)

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
          

3. Defining the CNN Architecture in PyTorch

We will construct a modular CNN class by subclassing `torch.nn.Module`. Our architecture will feature two convolutional blocks followed by fully connected classification layers.

              
            
C

            
            import torch.nn as nn
import torch.nn.functional as F

class ImageClassifierCNN(nn.Module):
    def __init__(self):
        super(ImageClassifierCNN, self).__init__()
        # Block 1: Input 3 channels (RGB), Output 32 filters, 3x3 kernel
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Block 2: Input 32 channels, Output 64 filters, 3x3 kernel
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        
        # Fully connected classification layers
        # CIFAR-10 images are 32x32 -> after two 2x2 max pools, spatial size is 8x8
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)  # 10 output classes
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        # First Conv -> ReLU -> MaxPool
        x = self.pool(F.relu(self.conv1(x)))
        # Second Conv -> ReLU -> MaxPool
        x = self.pool(F.relu(self.conv2(x)))
        
        # Flatten tensor for fully connected layers
        x = x.view(-1, 64 * 8 * 8)
        
        # Fully connected layers with dropout
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

model = ImageClassifierCNN()
print(model)
          

4. Configuring Loss, Optimizer, and Training Loop

To optimize our network weights, we use **Cross-Entropy Loss** (standard for multi-class classification) and the **Adam optimizer**.

              
            
C

            
            import torch.optim as optim

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Check for GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

epochs = 5
print(f"Training model on device: {device}")

for epoch in range(epochs):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        # 1. Zero gradients
        optimizer.zero_grad()
        
        # 2. Forward pass
        outputs = model(inputs)
        
        # 3. Calculate loss
        loss = criterion(outputs, labels)
        
        # 4. Backward pass and optimization step
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 200 == 199:    # Print every 200 mini-batches
            print(f"[Epoch {epoch + 1}, Batch {i + 1}] loss: {running_loss / 200:.3f}")
            running_loss = 0.0

print("Finished Training")
          

5. Evaluating Model Accuracy on Test Data

Once training is complete, we evaluate our trained CNN against unseen test images, calculating overall predictive accuracy.

              
            
C

            
            correct = 0
total = 0

# Disable gradient calculation for evaluation
model.eval()
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy of the network on the 10,000 test images: {100 * correct / total:.2f}%")
          

Conclusion and Best Practices

Data Augmentation: Incorporate random cropping, horizontal flipping, and color jittering during training to improve model generalization and prevent overfitting.
Batch Normalization: Insert `nn.BatchNorm2d` layers after convolutions to stabilize training dynamics and allow higher learning rates.
Leverage Pretrained Models: For complex production tasks, fine-tune established architectures like ResNet or EfficientNet from `torchvision.models` rather than training from scratch.

By understanding convolutional operations, tensor transformations, and modular architecture design in PyTorch, you are fully equipped to tackle advanced computer vision challenges.