Pytorch - storage and loading model

Interview questions ：

• PyTorch Of state_dict What does it contain ？

• PyTorch There are several ways to save models ,checkpoint How is it different from other ways , What do you usually keep ？

SAVING AND LOADING MODELS FOR INFERENCE IN PYTORCH

Two ways of preservation ：

1. state_dict,torch.nn.modules.module,Module class , Is the parent of multiple classes , Example layer 、 The optimizer etc.
   • state_dict function , Storage parameters and buffers, for example , The normalized value of the batch is buffers
2. All models

Net Inherited from Module,__init__ Initialization layer ,forward Connect layers , Input x, Instantiation net = Net()

Call optimizer optim.SGD, The first 1 The first parameter is the parameter of the model ,net.parameters() function , Include current and child module Parameters of

torch.save(net.state_dict(), PATH), With name 、epoch、train loss、eval loss, Save only parameters , There is no structure to save the model ( chart )

about Net Instantiation , call load_state_dict() function , hold dict Import in , Use torch.load(PATH)

preservation ：

• save -> state_dict
• load -> load_state_dict

call eval(), take training Set to False, The gradient will not be saved , Will also require_grad Set to false, Use reasoning patterns at the same time , for example Dropout、BN layer

torch.save(net, PATH), Keep the graph structure and parameters directly , Just call it directly ,torch.load(PATH)

import torch
import torch.nn as nn
import torch.optim as optim

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()
print(net)

optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# Specify a path
PATH = "state_dict_model.pt"

# Save
torch.save(net.state_dict(), PATH)

# Load
model = Net()
model.load_state_dict(torch.load(PATH))
model.eval()

# Specify a path
PATH = "entire_model.pt"

# Save
torch.save(net, PATH)

# Load
model = torch.load(PATH)
model.eval()

SAVING AND LOADING A GENERAL CHECKPOINT IN PYTORCH

Save and load general checkpoint

checkpoint preservation , call torch.save(), When epoch % 5 == 0 when , call torch.save(dict, PATH)

Common parameters ：epoch、model_state_dict、optimizer_state_dict、loss, During training , A very important amount of information

torch.load(PATH) load checkpoint, To assign a value

• model.load_state_dict()
• optimizer.load_state_dict()
• epoch
• loss

During training , Try to press checkpoint Way,

# Additional information
EPOCH = 5
PATH = "model.pt"
LOSS = 0.4

torch.save({
    
            'epoch': EPOCH,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS,
            }, PATH)
            
model = Net()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

SAVING AND LOADING MULTIPLE MODELS IN ONE FILE USING PYTORCH

Save and load multiple models in one file

With saving a single model checkpoint similar , Put the parameters of multiple models into a big dictionary , Load together again , To deal with

PATH = "model.pt"

torch.save({
    
            'modelA_state_dict': netA.state_dict(),
            'modelB_state_dict': netB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            }, PATH)
            
modelA = Net()
modelB = Net()
optimModelA = optim.SGD(modelA.parameters(), lr=0.001, momentum=0.9)
optimModelB = optim.SGD(modelB.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

Use docker Container creation environment ：

seaborn：https://seaborn.pydata.org/

Common software ：

Free of charge GPU resources ：Colaboratory