test.py

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

from model import ModernLeNet5
if("__main__"==__name__):
    batch_size = 64
    num_classes = 10
    learning_rate = 0.001
    num_epochs = 10
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.Compose([
                                                transforms.Resize((32, 32)),
                                                transforms.ToTensor(),
                                                transforms.Normalize(mean = (0.1307), std=(0.3081))
                                                ]), download=True)
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size,
                                              shuffle=True)  # Only load data as necessary

    model = ModernLeNet5().to(device)
    model.load_state_dict(torch.load("modern_lenet5.pth"))
    model.eval() # Set model to evaluation mode

    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        accuracy = 100 * correct / total
        print(f'Accuracy of the model on the test set: {accuracy:.2f}%')