This notebook demonstrates how to save and load models with PyTorch. This is important because we often want to load previously trained models to use in making predictions or to continue training on new data.

%matplotlib inline

import matplotlib.pyplot as plt
import numpy as np

import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms

# Define a transform to normalize the data
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/',
train=True,
transform=transform)
batch_size=64,
shuffle=True)

testset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/',
train=False,
transform=transform)
batch_size=64,
shuffle=True)


Helper function to show images.

def imshow(image, ax=None, title=None, normalize=True):
if ax is None:
fig, ax = plt.subplots()
image = image.numpy().transpose((1, 2, 0))

if normalize:
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
image = std * image + mean
image = np.clip(image, 0, 1)

ax.imshow(image)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.tick_params(axis='both', length=0)
ax.set_xticklabels('')
ax.set_yticklabels('')

return ax


Here we can see one of the images.

image, label = next(iter(trainloader))
imshow(image[0,:]);


# Train a network

To make things more concise here, I moved the model architecture and training code from the last part to a file called fc_model. Importing this, we can easily create a fully-connected network with fc_model.Network, and train the network using fc_model.train. I’ll use this model (once it’s trained) to demonstrate how we can save and load models.

import torch
from torch import nn
import torch.nn.functional as F

class Network(nn.Module):
def __init__(self, input_size, output_size, hidden_layers, drop_p=0.5):
super().__init__()

# Input to a hidden layer
self.hidden_layers = nn.ModuleList([nn.Linear(input_size, hidden_layers[0])])

# Add a variable number of more hidden layers
layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:])
self.hidden_layers.extend([nn.Linear(h1, h2) for h1, h2 in layer_sizes])

self.output = nn.Linear(hidden_layers[-1], output_size)

self.dropout = nn.Dropout(p=drop_p)

def forward(self, x):
for each in self.hidden_layers:
x = F.relu(each(x))
x = self.dropout(x)
x = self.output(x)

return F.log_softmax(x, dim=1)

def validation(model, testloader, criterion):
accuracy, test_loss = 0, 0

images = images.resize_(images.size()[0], 784)

output = model.forward(images)
test_loss += criterion(output, labels).item()

## Calculating the accuracy
# Model's output is log-softmax, take exponential to get the probabilities
ps = torch.exp(output)

# Class with highest probability is our predicted class, compare with true label
equality = (labels.data == ps.max(1)[1])

# Accuracy is number of correct predictions divided by all predictions,
#      so just take the mean
accuracy += equality.type_as(torch.FloatTensor()).mean()

return test_loss, accuracy

def train(model, trainloader, testloader, criterion, optimizer, epochs=5, print_every=40):

running_loss = 0
print('Epoch\tStep\tTraining Loss\tTest Loss\tTest Accuracy')
for e in range(epochs):
# Reset Step Counter
step = 0

# Model in training mode, dropout is on
model.train()

# Flatten images into a 784 long vector
images.resize_(images.size()[0], 784)

output = model.forward(images)
loss = criterion(output, labels)
loss.backward()
optimizer.step()

running_loss += loss.item()

if step % print_every == 0:
# Model in inference mode, dropout is off
model.eval()

# Turn off gradients for validation, will speed up inference
test_loss, accuracy = validation(model, testloader, criterion)

print("{:3}/{}\t{:4}\t{:13.3f}\t{:9.3f}\t{:13.3f}"
.format(e+1, epochs,
step,

running_loss = 0

# Make sure dropout and grads are on for training
model.train()

step += 1


Create the network, define the criterion and optimizer, and train

model = Network(784, 10, [512, 256, 128])
criterion = nn.NLLLoss()

Epoch	Step	Training Loss	Test Loss	Test Accuracy
1/3	   0	        0.002	    2.282	        0.178
1/3	 392	        0.362	    0.540	        0.799
1/3	 784	        0.252	    0.495	        0.823
2/3	   0	        0.091	    0.494	        0.811
2/3	 392	        0.229	    0.468	        0.830
2/3	 784	        0.220	    0.458	        0.837
3/3	   0	        0.085	    0.447	        0.842
3/3	 392	        0.209	    0.427	        0.842
3/3	 784	        0.205	    0.417	        0.848


As you can imagine, it’s impractical to train a network every time you need to use it. Instead, we can save trained networks then load them later to train more or use them for predictions.

The parameters for PyTorch networks are stored in a model’s state_dict. We can see the state dict contains the weight and bias matrices for each of our layers.

print("Our model: \n\n", model, '\n')
print("The state dict keys:\n")

for key in list(model.state_dict().keys()):
print(key)

Our model:

Network(
(hidden_layers): ModuleList(
(0): Linear(in_features=784, out_features=512, bias=True)
(1): Linear(in_features=512, out_features=256, bias=True)
(2): Linear(in_features=256, out_features=128, bias=True)
)
(output): Linear(in_features=128, out_features=10, bias=True)
(dropout): Dropout(p=0.5, inplace=False)
)

The state dict keys:

hidden_layers.0.weight
hidden_layers.0.bias
hidden_layers.1.weight
hidden_layers.1.bias
hidden_layers.2.weight
hidden_layers.2.bias
output.weight
output.bias


The simplest thing to do is simply save the state dict with torch.save. For example, we can save it to a file 'checkpoint.pth'.

torch.save(model.state_dict(),
'saving-models/checkpoint.pth')


Note that the file is relatively large at 2.3 MB.

%%bash
cd saving-models
ls -l

total 2224
-rw-r--r-- 1 ryan ryan 2271185 May 29 06:51 checkpoint.pth
-rw-r--r-- 1 ryan ryan    2603 May 28 06:07 saving-models_7_0.png


Then we can load the state dict with torch.load.

state_dict = torch.load('saving-models/checkpoint.pth')

for key in list(model.state_dict().keys()):
print(key)

hidden_layers.0.weight
hidden_layers.0.bias
hidden_layers.1.weight
hidden_layers.1.bias
hidden_layers.2.weight
hidden_layers.2.bias
output.weight
output.bias


And to load the state dict in to the network, you do model.load_state_dict(state_dict).

model.load_state_dict(state_dict)

<All keys matched successfully>


Seems pretty straightforward, but as usual it’s a bit more complicated. Loading the state dict works only if the model architecture is exactly the same as the checkpoint architecture. If you create a model with a different architecture, this fails:

# Try this
model = Network(784, 10, [400, 200, 100])
# This will throw an error because the tensor sizes are wrong!

---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-14-796f2ae97f9b> in <module>
2 model = Network(784, 10, [400, 200, 100])
3 # This will throw an error because the tensor sizes are wrong!

845         if len(error_msgs) > 0:
--> 847                                self.__class__.__name__, "\n\t".join(error_msgs)))
848         return _IncompatibleKeys(missing_keys, unexpected_keys)
849

size mismatch for hidden_layers.0.weight: copying a param with shape torch.Size([512, 784]) from checkpoint, the shape in current model is torch.Size([400, 784]).
size mismatch for hidden_layers.0.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([400]).
size mismatch for hidden_layers.1.weight: copying a param with shape torch.Size([256, 512]) from checkpoint, the shape in current model is torch.Size([200, 400]).
size mismatch for hidden_layers.1.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([200]).
size mismatch for hidden_layers.2.weight: copying a param with shape torch.Size([128, 256]) from checkpoint, the shape in current model is torch.Size([100, 200]).
size mismatch for hidden_layers.2.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([100]).
size mismatch for output.weight: copying a param with shape torch.Size([10, 128]) from checkpoint, the shape in current model is torch.Size([10, 100]).


This means we need to rebuild the model exactly as it was when trained. Information about the model architecture needs to be saved in the checkpoint, along with the state dict. To do this, you build a dictionary with all the information you need to compeletely rebuild the model.

checkpoint = {'input_size': 784,
'output_size': 10,
'hidden_layers': [each.out_features for each in model.hidden_layers],
'state_dict': model.state_dict()}

torch.save(checkpoint, 'saving-models/checkpoint.pth')


Now the checkpoint has all the necessary information to rebuild the trained model. You can easily make that a function if you want. Similarly, we can write a function to load checkpoints.

def load_checkpoint(filepath):

model = load_checkpoint('saving-models/checkpoint.pth')