PyTorch official Tutorials

Follow the of PyTorch's official Tutorials code to understand your slightly changed code and add comments. The jupyter notebook used by the IDE

Link: SAVE AND LOAD THE MODEL

Saving and loading models

Save the model state by saving, loading and predicting the model

import torch
import torch.onnx as onnx
import torchvision.models as models

Saving and loading model weights

The pytorch model stores the learned parameters in the internal state dictionary, which is called state_dict. These can be through torch Save method to save:

model=models.vgg16(pretrained=True)
torch.save(model.state_dict(),'model_weights.pth')

When loading model weights, you need to create an instance of the same model first, and then use load_ state_ The dict () method to load the parameters

model=models.vgg16() #There is no special declaration that pre trained = true, so the default weight will not be loaded
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

Ensure that model. is called before prediction. In the eval () method, the dropout and batch normalization layers are set to the evaluation mode Failure to do so will result in inconsistent predictions

Save model parameters with structure

When loading model weights, you need to instantiate the model class first, because the class defines the structure of the network We may want to save the class structure together with the model. At this time, we just need to pass model (not model.state_dict()) into the save function:

torch.save(model,'model.pth')

You can load the model as follows:

model=torch.load('model.pth')

In this way, python is used to serialize the model pickle Module, so it depends on the definition of the actual class available when loading the model

This sentence is not well understood

Export model to ONNX

Pytorch also supports raw? ONNX export However, considering the dynamic graph of pytorch itself, the export process must traverse the graph to produce a continuous ONNX model For this reason, a properly sized test variable should be passed into the export routine (in this case, a correctly sized zero tensor is used):

input_image=torch.zeros((1,3,244,244))
onnx.export(model,input_image,'model.onnx')

Many things can be done with ONNX model, including predicting in different languages on different platforms and recommending access ONNX tutorial
congratulations! You have completed the Getting Started tutorial for pytorch! You can now access quick start Quickly review the entire tutorial

This is the end of the introductory tutorial, and the subsequent advanced tutorials will not be translated into Chinese. Pay attention to the code comments as much as possible