This project outlines the process of training a Convolutional Neural Network (CNN) for image classification using PyTorch. The workflow includes data loading and preprocessing, model definition, training, evaluation, and testing.
The dataset used is MNIST—a collection of grayscale images of handwritten digits (0–9). All images are converted into PyTorch tensors for processing. The CNN is designed for grayscale inputs and is trained using a loss function and optimizer to improve classification accuracy. The training loop includes validation steps to monitor performance.
Note: Selected code snippets are included below to highlight key parts of the project. The complete Jupyter Notebook, including all code and outputs, is available using the link at bottom of the page.
"""Connecting to GPU"""
""" Description:
This module attempts to connect to the GPU.
- Checks PyTorch version
- Checks CUDA version and if cuDNN is enabled
- Returns the GPU name, if successful
"""
import torch
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.enabled)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))
device = "cuda:0" if torch.cuda.is_available() else "cpu"
"""Define CNN Architecture"""
""" Description:
This module defines a Convolutional Neural Network (CNN) for the MNIST dataset.
- Implements a neural network with three convolutional layers, each followed by max pooling
- Includes fully connected layers with dropout for classification
- Configures input and output dimensions for MNIST (28x28 grayscale images, 10 classes)
- Calculates flattened feature size dynamically for the fully connected layers
- Uses ReLU activations and dropout for regularization
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class NeuralNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=5, padding=2)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5, padding=2)
self.pool2 = nn.MaxPool2d(2, 2)
self.conv3 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, padding=1)
self.pool3 = nn.MaxPool2d(2, 2)
self._flattened_features = self._get_conv_output_size()
self.fc1 = nn.Linear(in_features=self._flattened_features, out_features=1024)
self.drop1 = nn.Dropout(p=0.3)
self.fc2 = nn.Linear(in_features=1024, out_features=512)
self.drop2 = nn.Dropout(p=0.3)
self.out = nn.Linear(in_features=512, out_features=10)
def _get_conv_output_size(self):
dummy_input = torch.zeros(1, 1, 28, 28)
dummy_input = self.conv1(dummy_input)
dummy_input = self.pool1(dummy_input)
dummy_input = self.conv2(dummy_input)
dummy_input = self.pool2(dummy_input)
dummy_input = self.conv3(dummy_input)
dummy_input = self.pool3(dummy_input)
return int(torch.flatten(dummy_input, 1).size(1))
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool1(x)
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = F.relu(self.conv3(x))
x = self.pool3(x)
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = self.drop1(x)
x = F.relu(self.fc2(x))
x = self.drop2(x)
x = self.out(x)
return x
"""CNN Training and Evaluation"""
""" Description:
This module trains and evaluates a Convolutional Neural Network (CNN) on the MNIST dataset.
- Trains the model over multiple epochs
- Validates the model after each epoch
- Collects predictions from the test set
- Computes accuracy, classification report, and confusion matrix
"""
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
num_epochs = 10
for epoch_index in range(num_epochs):
print(f'Epoch: {epoch_index + 1}\n')
train_one_epoch()
validate_one_epoch()
print('Finished Training')
predicted_labels = []
true_labels = []
for images, labels in testloader:
images = images.to(device)
labels = labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)
predicted_labels.extend(predicted.cpu().numpy())
true_labels.extend(labels.cpu().numpy())
accuracy = accuracy_score(true_labels, predicted_labels)
print("Accuracy:", accuracy)
class_report = classification_report(true_labels, predicted_labels, target_names=classes)
conf_matrix = confusion_matrix(true_labels, predicted_labels)
This project features a deep learning model built with PyTorch to recognize handwritten digits from the MNIST dataset, achieving an accuracy of 99.33%. I created a neural network to classify images of digits (0–9), using efficient data processing and GPU acceleration. The model was trained and tested , and its performance was evaluated with detailed metrics.
Visit the Deep Learning with Pytorch Project on GitHub