Loading and running a trained model with PyTorch

This notebook walks you through how to load and run a trained model using PyTorch. You can run each cell one by one as you work through the content, or run all cells and read through all output at the same time.

In the previous notebook, torch-use-gpu.ipynb, we performed the following tasks:

Imported Torch libraries
Listed available GPUs.
Checked that GPUs are enabled.
Assigned a GPU device and retrieved the device name.
Loaded vectors, matrices, and data onto a GPU.
Loaded a neural network model onto a GPU.
Trained the neural network model.

Make sure that you have run the torch-use-gpu.ipynb notebook file before working through this notebook.

In this notebook, you will:

Import libraries and GPUs into your notebook environment.
Load a model and training data onto a GPU.
Run a model on a GPU.
Review model results.

Run the model with your trained weights¶

The first thing to to is set up the neural network model to run using the weights you created when you trained the model in the torch-use-gpu.ipynb file.

Set up your environment¶

Import the torch and torchvision libraries you need in order to work with PyTorch, and ensure that your GPU resources are visible in your notebook server.

!pip install torchvision==0.9.1
!pip install tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)  # let's see what device we got

Load the model and trained weights onto the GPU¶

Run the following commands to set up a simple neural network model.

# Getting set to import our trained model.

# batch size of 1 so we can look at one image at time.
batch_size = 1


class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 784)
        self.fc2 = nn.Linear(784, 10)

    def forward(self, x):
        x = x.view(batch_size, -1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        output = F.softmax(x, dim=1)
        return output

Then, load the model onto the GPU along with the trained weights from the torch-use-gpu.ipynb workbook.

model = SimpleNet().to( device )
model.load_state_dict( torch.load("mnist_fashion_SimpleNet.pt") )

model.eval()

# using our loader as before - will have that same batch size of 1.

test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
        ])


testset = datasets.FashionMNIST('./data', train=False,
                                transform=test_transform)
test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                          shuffle=False, num_workers=2)

testing_iterator = iter(test_loader)  # create an iterator of our loader

# A dictionary to map our class numbers to their items.
labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}

def get_probs(model, x, y, device):
    """ Function to return the current probabilities of the classes given the model """
    model = model.to(device)
    x = x.to(device)
    y = y.to(device)
    
    with torch.no_grad():
        output = model(x)  # model classifies the input
        output = output.squeeze()
        
    # Return the list of probabilites and the probability of ground truth class
    return output, output[y]

Run the model and check the results¶

With the model and trained weights imported, we can now run our model and see how well it performs.

Run the below cell repeatedly to see if the model predicts the selected image correctly.

image, label = testing_iterator.next()  # get one datapoint
image = image.squeeze(0)  # 'squeeze' out the tensor from the outer list.
print("image shape is", image.shape)

probs, gt_prob = get_probs(model, image, label, device)  # run model prediction


# plot our image
plt.imshow(image.view(28,28), cmap="gray") 
plt.show()

# what did the model predict?
max_label = torch.argmax(probs)
print("Model predicted:", labels_map[max_label.item()], "with confidence % of", probs[max_label.item()].item())
print("Correct label is:", labels_map[label.item()])

The model predicts the category that the photographed object should be in, and lists its confidence percentage beside the prediction.