image_classifier.py

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

#Define steps to process image
#   transforms.ToTensor converts the image to a pytorch tensor
#   transforms.Normalize, normalizes the data from being [0.0,1.0] to [-1.0,1.0]
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,),(0.5,))])

#Get training data
trainset = torchvision.datasets.MNIST(root="./data",train=True,download=True,transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,shuffle=True)

#Get test data
testset = torchvision.datasets.MNIST(root="./data",train=False,download=True,transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64,shuffle=True)

# ---------------------------------------------------------------
# 2. DEFINE THE NEURAL NETWORK
# ---------------------------------------------------------------
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        #Convolutional layers
        self.conv1 = nn.Conv2d(1, 32, 3, 1)   #Input 1x28x28 -> Output 32x26x26
        self.conv2 = nn.Conv2d(32, 64, 3, 1)  #Input 32x26x26 -> Output 64x24x24
        
        #Fully connected layers
        #After max pooling, feature maps are 64x12x12 -> flattened to 9216
        self.fc1 = nn.Linear(9216, 128)  #Combines features into 128 high-level features
        self.fc2 = nn.Linear(128, 10)    #Output layer: 10 classes (digits 0-9)

    def forward(self,x):
        #Apply convolution + ReLU
        #Introduces non-linearity to convo layers by setting negative data to be zero
        x = F.relu(self.conv1(x)) 
        x = F.relu(self.conv2(x))

        #Halves the image width and height
        x = F.max_pool2d(x,2) 
        #It is 2 because we have to kernels or 2 convolution layers

        #Flatten 64x12x12 feature maps to 1D vector (9216)
        x = torch.flatten(x,1)
        
        #Fully connected layer + ReLU
        #Makes all negative elements in 9216 element vector 0
        x = F.relu(self.fc1(x))
        
        # Output layer (logits for 10 classes)
        x = self.fc2(x)

        #Returns that final 10 element vector
        return x

# Create the network instance of nn 
net = Net()

# ---------------------------------------------------------------
# 3. LOSS FUNCTION AND OPTIMIZER
# ---------------------------------------------------------------
criterion = nn.CrossEntropyLoss()             # Loss function for classification
optimizer = optim.Adam(net.parameters(), lr=0.001)  # Optimizer for training

# ---------------------------------------------------------------
# 4. TRAINING LOOP
# ---------------------------------------------------------------
print("Training started...\n")

#Train for 5 epochs
for epoch in range(5):
    running_loss =0.0

    #Iterate over every batch in the training dataset
    for images, labels in trainloader:
        optimizer.zero_grad()           # Reset gradients
        outputs = net(images)           # Forward pass
        loss = criterion(outputs, labels)  # Compute loss
        loss.backward()                 # Backpropagate
        optimizer.step()                # Update weights

        running_loss += loss.item()

    print(f"Epoch [{epoch+1}/5] - Loss: {running_loss/len(trainloader):.4f}")

print("\nTraining complete!")

# ---------------------------------------------------------------
# 5. TESTING / EVALUATION
# ---------------------------------------------------------------
correct = 0
total = 0

#Turns of gradient not needed for testing
with torch.no_grad():
    for images,labels in testloader:
        #Runs images through net foward function
        outputs = net(images)

        #Finds highest value in the each output, there are 10 scores in total
        _, predicted = torch.max(outputs.data, 1)

        #Total number of samples tested
        total += labels.size(0)

        #Checks each value in the tensor and sees if correct
        #If correct true(1) else wrong(0)sums them all up
        #However still tensor example tensor(48), .item() get the item in the tensor
        correct += (predicted == labels).sum().item()

print(f"Test Accuracy: {100 * correct / total:.2f}%")

# ---------------------------------------------------------------
# 6. VISUALIZE PREDICTIONS
# ---------------------------------------------------------------
dataiter = iter(testloader)
images, labels = next(dataiter)

# Get predictions
outputs = net(images)
_, preds = torch.max(outputs, 1)

# Unnormalize the first image (undo [-1,1] → [0,1])
images = images / 2 + 0.5
npimg = images[0].numpy().squeeze()

plt.imshow(npimg, cmap='gray')
plt.title(f"Predicted: {preds[0].item()} | Actual: {labels[0].item()}")
plt.show()