Identifying handwritten digits using Logistic Regression in PyTorch?

This tutorial demonstrates how to build a Convolutional Neural Network (CNN) using PyTorch to classify handwritten digits from the MNIST dataset. We'll create a CNN model and train it to achieve high accuracy on digit recognition.

The MNIST dataset contains 70,000 labeled 28×28 pixel grayscale images of handwritten digits (0-9), with 60,000 training images and 10,000 test images.

Installation and Setup

First, install the required libraries ?

pip install torch torchvision matplotlib

Import Libraries

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt

Setting Hyperparameters

Define the training parameters and random seed for reproducible results ?

# Training parameters
n_epochs = 3
batch_size_train = 64
batch_size_test = 1000
learning_rate = 0.01
momentum = 0.5
log_interval = 10

# Set random seed for reproducibility
random_seed = 1
torch.backends.cudnn.enabled = False
torch.manual_seed(random_seed)

<torch._C.Generator at 0x2048010c690>

Loading the MNIST Dataset

Load the MNIST dataset with normalization using the dataset's mean (0.1307) and standard deviation (0.3081) ?

# Data loaders for training and testing
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data/', train=True, download=True,
        transform=torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.1307,), (0.3081,))
        ])
    ),
    batch_size=batch_size_train, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data/', train=False, download=True,
        transform=torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.1307,), (0.3081,))
        ])
    ),
    batch_size=batch_size_test, shuffle=True)

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Processing...
Done!

Examining the Data

Let's check the shape of our test data batch ?

examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
print(f"Data shape: {example_data.shape}")

Data shape: torch.Size([1000, 1, 28, 28])

The output shows we have 1000 examples of 28×28 pixel grayscale images (1 channel).

Building the CNN Model

Create a CNN with two convolutional layers followed by two fully connected layers ?

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # First convolutional layer: 1 input channel, 10 output channels, 5x5 kernel
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        # Second convolutional layer: 10 input channels, 20 output channels, 5x5 kernel
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        # Dropout layer for regularization
        self.conv2_drop = nn.Dropout2d()
        # Fully connected layers
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)  # 10 output classes for digits 0-9

    def forward(self, x):
        # First conv layer + ReLU + max pooling
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        # Second conv layer + dropout + ReLU + max pooling
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        # Flatten for fully connected layers
        x = x.view(-1, 320)
        # First fully connected layer + ReLU
        x = F.relu(self.fc1(x))
        # Dropout for regularization
        x = F.dropout(x, training=self.training)
        # Output layer with log softmax
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# Initialize network and optimizer
network = Net()
optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)

Training and Testing Functions

Define functions for training and testing the model ?

def train(epoch):
    network.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = network(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        
        if batch_idx % log_interval == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

def test():
    network.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = network(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            pred = output.data.max(1, keepdim=True)[1]
            correct += pred.eq(target.data.view_as(pred)).sum()
    
    test_loss /= len(test_loader.dataset)
    print(f'\nTest set: Avg. loss: {test_loss:.4f}, '
          f'Accuracy: {correct}/{len(test_loader.dataset)} '
          f'({100. * correct / len(test_loader.dataset):.0f}%)\n')
    return test_loss

Training the Model

Train the model for the specified number of epochs ?

# Test initial performance with random weights
test()

# Train for specified epochs
for epoch in range(1, n_epochs + 1):
    train(epoch)
    test()

Test set: Avg. loss: 2.3048, Accuracy: 1063/10000 (10%)

Train Epoch: 1 [0/60000 (0%)]	Loss: 2.294911
Train Epoch: 1 [640/60000 (1%)]	Loss: 2.314225
...
Train Epoch: 3 [59520/60000 (99%)]	Loss: 0.318569

Test set: Avg. loss: 0.0912, Accuracy: 9716/10000 (97%)

Model Performance

After just 3 epochs of training, the model achieved 97% accuracy on the test set, improving from the initial 10% accuracy with random weights.

Making Predictions

Test the model on some example images ?

with torch.no_grad():
    output = network(example_data)
    predictions = output.data.max(1, keepdim=True)[1]
    
    # Show first 6 predictions
    print("Predictions for first 6 test images:")
    for i in range(6):
        print(f"Image {i+1}: Predicted = {predictions[i].item()}, "
              f"Actual = {example_targets[i].item()}")

Predictions for first 6 test images:
Image 1: Predicted = 7, Actual = 7
Image 2: Predicted = 2, Actual = 2
Image 3: Predicted = 1, Actual = 1
Image 4: Predicted = 0, Actual = 0
Image 5: Predicted = 4, Actual = 4
Image 6: Predicted = 1, Actual = 1

Conclusion

This CNN successfully learned to classify handwritten digits with 97% accuracy using PyTorch. The model combines convolutional layers for feature extraction with fully connected layers for classification, demonstrating the effectiveness of deep learning for image recognition tasks.