How to Build and Train a Deep Learning Image Classifier in PyTorch (CIFAR-10)

📅 April 21, 2025 📂 Machine Learning, Featured

Objective:

In this tutorial, we will walk through the process of training an image classifier using PyTorch. We will focus on understanding PyTorch, exploring the CIFAR-10 dataset, and building and training a Convolutional Neural Network (CNN) to classify images.

Introduction:

PyTorch is an open-source machine learning library that has gained immense popularity in recent years. It provides a flexible, easy-to-use framework for building deep learning models and performing complex computations. What makes PyTorch unique is its dynamic computation graph, which allows for easy debugging and experimentation.

In this tutorial, we will use the CIFAR-10 dataset—a popular dataset used for image classification tasks. The CIFAR-10 dataset consists of 60,000 32×32 color images in 10 classes, such as dogs, cats, and airplanes. It’s a great starting point for anyone who is just beginning with image classification.

Section 1: What is PyTorch?

PyTorch is a Python-based library for deep learning that emphasizes flexibility and speed. It provides dynamic computation graphs, meaning that the graph is defined on-the-fly during runtime, which makes debugging and modification easier compared to other frameworks. You can modify your model architecture and operations dynamically without needing to recompile the entire model.

Key features of PyTorch:

Dynamic Computation Graphs: Unlike static computation graphs, dynamic graphs are created and modified during the execution of your model, making debugging and modification easy.
Autograd: PyTorch’s autograd system automatically computes gradients, saving time when working with optimization algorithms like gradient descent.
Torchvision: PyTorch provides torchvision, a library designed to make working with image data easier, offering datasets, pre-trained models, and common image transformations.

These features make PyTorch highly popular for deep learning research and production applications.

Section 2: The CIFAR-10 Dataset

The CIFAR-10 dataset is a widely used benchmark for image classification. It contains 60,000 color images, each 32×32 pixels, across 10 classes. The dataset is split into:

50,000 training images
10,000 test images

The 10 classes include:

Airplane
Automobile
Bird
Cat
Deer
Dog
Frog
Horse
Ship
Truck

For this tutorial, we’ll use CIFAR-10 for training a CNN model. We will also explore the preprocessing steps, such as normalizing the images and converting them into tensors for efficient processing by the model.

Section 3: Setting Up PyTorch and Installing Dependencies

To begin, you need to install PyTorch and torchvision. Use the following command to install them:

pip install torch torchvision

Once installed, you can import the necessary libraries in your code:

import torch
import torchvision
import torchvision.transforms as transforms

We will also use torchvision for loading the CIFAR-10 dataset and transforming the images.

Section 4: Preparing the CIFAR-10 Dataset for Training

Now that we have installed PyTorch and torchvision, let’s prepare the CIFAR-10 dataset. We’ll load the dataset and apply transformations such as converting images to tensors and normalizing the data.

pythonCopyEdit# Transformations to normalize the dataset
transform = transforms.Compose([
    transforms.ToTensor(),  # Converts images to PyTorch tensors
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalizes the images
])

# Downloading the CIFAR-10 training dataset
trainset =.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

What is happening here?

transforms.Compose(): This is used to chain together multiple transformations that will be applied to the dataset.
transforms.ToTensor(): This converts each image into a PyTorch tensor. PyTorch models work with tensors, and this transformation prepares the data for further operations.
transforms.Normalize(): This step normalizes the pixel values of the image, ensuring that each pixel has a mean of 0 and a standard deviation of 1. Normalization is important for model training because it prevents large differences in the scale of the data, making the training process more stable.

# Loading the dataset into a DataLoader
trainloader =.utils..DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

torch.utils.data.DataLoader(): This creates a DataLoader, which efficiently loads the dataset in batches. A batch is a subset of the data that is passed through the model at once. This speeds up the training process by allowing the model to process multiple samples simultaneously.
batch_size=4: Specifies that we will process 4 images at a time in each batch.
shuffle=True: Shuffling the data ensures that the model sees the data in a different order each time, which helps prevent overfitting.
num_workers=2: This tells the DataLoader to use 2 parallel processes to load the data, speeding up the data loading.

Section 5: Building the CNN Model

A Convolutional Neural Network (CNN) is a deep learning model particularly effective for image classification. CNNs consist of multiple layers, including convolutional layers, pooling layers, and fully connected layers. Let’s build a simple CNN for CIFAR-10.

import torch.nn as nn
import torch.optim as optim

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        
        # Convolutional layer: 3 input channels (RGB), 32 output channels, 3x3 kernel
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)  # Pooling layer with 2x2 kernel
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)  # Second convolutional layer
        self.fc1 = nn.Linear(64 * 8 * 8, 512)  # Fully connected layer
        self.fc2 = nn.Linear(512, 10)  # Output layer (10 classes)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))  # Apply Conv1 and MaxPooling
        x = self.pool(torch.relu(self.conv2(x)))  # Apply Conv2 and MaxPooling
        x = x.view(-1, 64 * 8 * 8)  # Flatten the output of Conv2
        x = torch.relu(self.fc1(x))  # Fully connected layer 1
        x = self.fc2(x)  # Output layer
        return x

# Initialize the model
model = CNN()

Explanation of the layers:

nn.Conv2d(3, 32, 3, padding=1): This creates the first convolutional layer. It takes 3 input channels (for RGB images), produces 32 output channels (feature maps), and uses a 3×3 kernel.
self.pool = nn.MaxPool2d(2, 2): This is the max-pooling layer, which reduces the spatial dimensions (height and width) of the image by a factor of 2.
self.fc1 = nn.Linear(64 * 8 * 8, 512): This is a fully connected layer. The input size is the flattened output from the convolutional layers (64 channels, each of size 8×8), and the output is 512 features.
self.fc2 = nn.Linear(512, 10): This is the final output layer with 10 neurons, one for each class in CIFAR-10.

Section 6: Training the Model

Now, we will train the model using the Stochastic Gradient Descent (SGD) optimizer and CrossEntropyLoss for multi-class classification.

# Loss function and optimizer
criterion = nn.CrossEntropyLoss()  # Loss function for classification
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)  # Optimizer

# Training loop
for epoch in range(10):  # Loop over the dataset multiple times (10 epochs)
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()  # Zero the parameter gradients

        # Forward pass
        outputs = model(inputs)
        loss = criterion(outputs, labels)  # Calculate the loss

        # Backward pass
        loss.backward()
        optimizer.step()  # Update weights

        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(trainloader)}")

Key components:

optimizer.zero_grad(): Clears the gradients of all optimized variables before the backward pass to ensure gradients are not accumulated from previous steps.
loss.backward(): Computes the gradient of the loss function with respect to the model parameters using backpropagation.
optimizer.step(): Updates the model parameters using the computed gradients.
running_loss: Tracks the average loss for the epoch to monitor training progress.

Section 7: Evaluating the Model

After training, we can evaluate the model’s performance on the test set.

# Testing the model
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)

correct = 0
total = 0

with torch.no_grad():  # Disable gradient tracking for inference
    for data in testloader:
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total}%')

Explanation:

torch.no_grad(): Disables gradient computation during inference (testing), which saves memory and computation time.
torch.max(outputs, 1): Finds the predicted class (the class with the highest score) for each image.
correct += (predicted == labels).sum().item(): Compares the predicted classes to the true labels and counts the number of correct predictions.

Section 8: Saving the Trained Model

After training, we can save the model so we don’t need to retrain it every time.

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

This saves the model’s parameters, allowing you to load and use the model later without retraining.

Conclusion:

In this tutorial, we’ve walked through the process of building and training a Convolutional Neural Network (CNN) using PyTorch and the CIFAR-10 dataset. We covered everything from loading the data to training the model, evaluating its performance, and saving it for later use.

This is just the beginning! There are many ways you can improve this model, including experimenting with different architectures, fine-tuning hyperparameters, and exploring advanced topics like data augmentation and transfer learning.