Building an Image Classifier with a Single-Layer Neural Network in PyTorch

by

A single-layer neural network, also known as a single-layer perceptron, represents the simplest form of neural networks. Comprising just one layer of neurons connected to both the input and output layers, this model can be effectively used for tasks such as image classification where the input layer receives image data, and the output layer produces class labels.

In this tutorial, you will learn how to build an image classifier using a single-layer neural network implemented in PyTorch. Specifically, you will explore:

  • How to prepare your data.
  • How to define the neural network model.
  • How to train and evaluate the classifier.

Let’s get started!

Overview

This tutorial consists of three main parts:

  1. Preparing the Dataset
  2. Building the Model Architecture
  3. Training the Model

Preparing the Dataset

This tutorial will utilize the CIFAR-10 dataset, which contains 60,000 color images of 32×32 pixels across 10 classes, with 6,000 images per class. The dataset includes categories like airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. CIFAR-10 is a well-known dataset in the realms of machine learning and computer vision, making it an ideal choice for your project.

You can easily import the CIFAR-10 dataset using PyTorch with the following code:

import torch
import torchvision
import torchvision.transforms as transforms

# Load the CIFAR-10 dataset
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())

The dataset will be downloaded automatically if it’s not already present in your specified root directory. To verify the number of training and testing samples, run:

print("Number of training samples: ", len(train_set))
print("Number of testing samples: ", len(test_set))

You should get the following output:

Number of training samples: 60000
Number of testing samples: 10000

Each sample consists of an image and a corresponding label. You can inspect the datatype and size of the first training sample by using:

print("Data type of the first training sample: ", train_set[0][0].type())
print("Size of the first training sample: ", train_set[0][0].size())

This will yield:

Data type of the first training sample: torch.FloatTensor
Size of the first training sample: torch.Size([3, 32, 32])

Here, the first image is represented as a 32×32 pixel RGB image (3 color channels).

Building the Model Architecture

Next, you’ll create a simple neural network, referred to as SimpleNet, using PyTorch’s nn.Module. This model will have one fully connected layer. The following code illustrates how to define this model:

import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(32*32*3, 100)  # Input to first layer with 100 hidden neurons
        self.fc2 = nn.Linear(100, num_classes)  # Second layer with output neurons equal to number of classes

    def forward(self, x):
        x = x.view(-1, 32*32*3)  # Flatten the input
        x = torch.relu(self.fc1(x))  # Apply ReLU activation
        x = self.fc2(x)  # Output layer
        return x

Next, instantiate the model:

# Instantiate the model
model = SimpleNet()

Training the Model

We will create two instances of PyTorch’s DataLoader class for both training and testing. For the training DataLoader, set the batch size to 64 and enable shuffling to prevent model overfitting.

Define the loss function and optimizer for training the model:

from torch.utils.data import DataLoader

# Load the data into PyTorch DataLoader
train_loader = DataLoader(dataset=train_set, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_set, batch_size=64, shuffle=False)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()  # Suitable for multi-class classification
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)  # Stochastic Gradient Descent

Now, let’s set up the training loop to train our model for a specified number of epochs. We will track the loss and accuracy during training:

num_epochs = 20
train_loss_history = []
train_acc_history = []
val_loss_history = []
val_acc_history = []

for epoch in range(num_epochs):
    train_loss = 0.0
    train_acc = 0.0

    # Set model to training mode
    model.train()
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)  
        loss = criterion(outputs, labels)  
        train_loss += loss.item()
        train_acc += (outputs.argmax(1) == labels).sum().item()  # Count correct predictions
        loss.backward()  # Backpropagation
        optimizer.step()  # Update model parameters

    # Calculate average training loss and accuracy
    train_loss /= len(train_loader)
    train_loss_history.append(train_loss)
    train_acc /= len(train_loader.dataset)
    train_acc_history.append(train_acc)

    print(f'Epoch: {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}')

print("Training completed!")

Upon training completion, you will observe progress reporting the training loss and accuracy.

Summary

In this tutorial, you learned how to build an image classifier using a single-layer neural network with PyTorch. Specifically, you covered:

  • How to prepare and load the CIFAR-10 dataset.
  • The construction of a simple neural network model.
  • The implementation of the training process with cross-entropy loss.

This base knowledge sets you up for more advanced deep learning tasks and applications using PyTorch. As you deepen your understanding, you’ll be better equipped to tackle diverse machine learning challenges.

Leave a Comment