Training an Image Classification Model using PyTorch

Training an Image Classification Model using PyTorch

Image classification is one of the most fundamental tasks in computer vision and deep learning. It involves training a model to categorize images into predefined classes or labels. PyTorch, a popular deep learning framework, provides a robust platform for building and training image classification models. In this blog, we’ll take you through the process of training an image classification model using PyTorch.

Prerequisites

Before we dive into the details, ensure you have the following prerequisites:

• Python: Make sure you have Python installed on your system.
• PyTorch: Install PyTorch by following the installation instructions on the official PyTorch website.
• GPU (optional): While not mandatory, using a GPU can significantly speed up training, especially for large models and datasets.

Dataset Selection

Selecting the right dataset is crucial for your image classification task. Common datasets include CIFAR-10, CIFAR-100, and ImageNet for general tasks, while more specific datasets like MNIST or Fashion MNIST are suitable for simpler tasks.

For this blog, we’ll use the CIFAR-10 dataset, which contains 60,000 32x32 color images in 10 different classes, with 6,000 images per class. You can easily access CIFAR-10 in PyTorch using the torchvision library.

Data Loading

To load and preprocess the dataset, you’ll use PyTorch’s DataLoader and transforms. The DataLoader class helps in efficient data loading, and transforms allow you to apply various data augmentation techniques like resizing, cropping, and normalization.

import torch
import torchvision
import torchvision.transforms as transforms

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

# Load CIFAR-10 dataset
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

Model Architecture

Selecting an appropriate model architecture is crucial for the success of your image classification task. For beginners, a simple Convolutional Neural Network (CNN) is a good starting point. You can create a basic CNN using PyTorch’s nn module:

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

Loss Function and Optimizer

For image classification, the common loss function used is the Cross-Entropy Loss, and the optimizer of choice is usually stochastic gradient descent (SGD). You can define these as follows:

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Training Loop

The training loop is where the magic happens. Here’s a basic structure for a training loop:

for epoch in range(2):  # Loop over the dataset multiple times
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # Get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # Zero the parameter gradients
        optimizer.zero_grad()

        # Forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # Print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:  # Print every 2000 mini-batches
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')

Evaluation

To evaluate your model, you can use a separate test dataset. Here’s a basic evaluation loop:

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')

Conclusion

Training an image classification model using PyTorch is a fundamental but crucial task in the field of deep learning and computer vision. This blog provided a step-by-step guide on data loading, model architecture, loss function, optimizer, training loop, and evaluation. With this foundation, you can start building more complex image classification models and explore various deep learning techniques to improve your model’s performance. Good luck with your deep learning journey!