train.py

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

from model import ModernLeNet5

if("__main__"==__name__):
    # Values as per LeNet5
    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")
    print("Using {} device".format(device))

    train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.Compose([
                                                transforms.Resize((32, 32)),
                                                transforms.ToTensor(),
                                                transforms.Normalize(mean = (0.1307), std=(0.3081))
                                                ]), download=True)

    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) # Only load data when being used


    model = ModernLeNet5().to(device)

    # Loss function
    cost = nn.CrossEntropyLoss()

    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    total_step = len(train_loader)
    for epoch in range(num_epochs):
        for step, (images, labels) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)

            # Forward pass
            outputs = model(images) # Retrieve outputs on our given images
            loss = cost(outputs, labels) # Test the loss on the outputs compared to the actual labels of the images
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if(((step+1) % 400) == 0):
                print("Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}".format(epoch + 1, num_epochs, step + 1, total_step, loss.item()))
    torch.save(model.state_dict(), "modern_lenet5.pth")