PyTorch官方Tutorials SAVE AND LOAD THE MODEL

PyTorch官方Tutorials

跟着PyTorch官方Tutorials码的，便于理解自己稍有改动代码并添加注释，IDE用的jupyter notebook

链接: SAVE AND LOAD THE MODEL

保存与加载模型

通过保存，加载，预测模型来保存模型状态

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

保存与加载模型权重

pytorch模型在内部状态字典保存习得的参数，这个字典叫做state_dict.这些可以通过torch.save方法来保存:

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

加载模型权重时,需要先创建一个同模型的实例,然后通过使用load_state_dict()方法来加载参数

model=models.vgg16() #这里没有特别声明 pretrained=True 因此不会加载默认权重
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)
  )
)

确保在预测前调用model.eval()方法来讲dropout和batch normalization层设置为评估模式.不这样做的话将会导致预测结果不一致.

将模型参数和结构一起保存

加载模型权重时,需要先实例化模型类,因为类定义了网络的结构.我们可能会想要和模型一起保存类的结构,此时只需将model(不是model.state_dict())传入保存函数:

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

可以这样加载模型:

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

这个方式在序列化模型时使用了python的pickle模块,因此它依赖于加载模型时可用的实际类的定义.

这句没咋理解

将模型导出到ONNX

pytorch也支持原始?ONNX导出.但是,考虑到pytorch本身的动态计算图,导出过程必须遍历计算图来产生一个持续的ONNX模型.出于这个原因,合适大小的测试变量应当传入导出例程(本例中使用一个正确大小的零张量):

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

用ONNX模型可以做很多事,包括在不同平台用不同语言进行预测,推荐访问ONNX教程
恭喜!您已经完成了pytorch的入门教程!现在可以访问quick start快速回顾整个教程

这里算是结束了入门教程,后续进阶教程就不做汉化了,注意点尽量在代码注释中体现