digit_classifier.py

'''
Module Name: digit_classifier.py
Author: Danish I.
Description: Using an MLP that has 2 hidden layers with 16 neurons each, classifies handwritten digits from 0-9 with roughly 90% accuracy.
'''
from neural_network import *
import torchvision
import json
import time

# Download and load
train_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    download=True,  
)

test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    download=True,
)

# Converts a label to a list where the label is set to 1 and all other values are set to 0. This is for computing the loss function.
def one_hot_encoding(label):
    output = [0.0 for _ in range(10)]
    output[label] = 1.0
    return output


network = MLP(784, [16, 16, 10])
epochs = 0

# Load previous weights
with open('model.json', 'r') as f:
    weights = json.load(f)
for p, w in zip(network.params(), weights):
    p.data = w

# Training loop
for epoch in range(epochs):
    correct = 0
    totalloss = 0
    start = time.time()
    for i in range(20000):
        # Process data
        image, label = train_dataset[i]
        image = [pixel / 255.0 for pixel in image.get_flattened_data()]

        # Forward pass
        ypred = network(image)
        loss = sum((yout - y_actual)**2 for yout, y_actual in zip(ypred, one_hot_encoding(label)))
        totalloss += loss.data

        # Backpropagation
        for p in network.params():
            p.grad = 0
        loss.backward()

        for p in network.params():
            p.data += -0.001 * p.grad

        # Calculate accuracy
        predicted_digit = max(range(10), key=lambda j: ypred[j].data) # the max is the models best prediction
        if predicted_digit == label:
            correct += 1

    elapsed = time.time() - start
    print(f"Epoch {epoch} average loss: {totalloss/20000:.4f}, accuracy: {correct/20000*100:.1f}%, time: {elapsed/3600:.2f}hrs")


# After training loop
weights = [p.data for p in network.params()]
with open('model.json', 'w') as f:
    json.dump(weights, f)

# Testing the model using the data from the testing data set
correct = 0
for i in range(10000):
    image, label = test_dataset[i]
    image = [pixel / 255.0 for pixel in image.get_flattened_data()]
    ypred = network(image)
    predicted_digit = max(range(10), key=lambda j: ypred[j].data)
    if predicted_digit == label:
        correct += 1

print(f"Test accuracy: {correct/10000*100:.1f}%")