17.初识Pytorch保存模型(model save)与加载模型(model load)

模型的保存有两种方式:一种是保存模型;另一种是保存模型的参数,将参数以字典的形式保存(官方推荐)。

There are two ways to save the model: one is to save the model; the other is to save the parameters of the model, and save the parameters in the form of a dictionary (official recommendation). 

code:

import torch
import torchvision

vgg16_false = torchvision.models.vgg16(pretrained=False)

vgg16_true = torchvision.models.vgg16(pretrained=True)
# method1 to save model
torch.save(vgg16_false, "vgg16_method1.pth")

# method2 to save parameters
torch.save(vgg16_true.state_dict(), "vgg16_method2.pth")

result
17.初识Pytorch保存模型(model save)与加载模型(model load)_第1张图片

加载模型:可以加载第一个以模型形式保存的文件;也可以加载第二个以模型的参数形式保存的文件,并且能把其转化成模型。

Load model: You can load the first file saved in the form of a model;
you can also load the second file saved in the form of model parameters, and convert it into a model. 

code:

import torch
import torchvision

# method1
model1 = torch.load("vgg16_method1.pth")
print(model1)

# method2
print("************************************")
model2 = torch.load("vgg16_method2.pth")
print(model2)

result:
17.初识Pytorch保存模型(model save)与加载模型(model load)_第2张图片17.初识Pytorch保存模型(model save)与加载模型(model load)_第3张图片17.初识Pytorch保存模型(model save)与加载模型(model load)_第4张图片17.初识Pytorch保存模型(model save)与加载模型(model load)_第5张图片

将保存参数字典的文件,加载成原本的模型.

Load the file that saves the parameter dictionary into the original model.

code:

import torch
import torchvision


# method2
print("************************************")
model2 = torch.load("vgg16_method2.pth")
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(model2)
print(vgg16)

result:
17.初识Pytorch保存模型(model save)与加载模型(model load)_第6张图片17.初识Pytorch保存模型(model save)与加载模型(model load)_第7张图片17.初识Pytorch保存模型(model save)与加载模型(model load)_第8张图片

如果是自己定义的模型,进行保存和加载,则需要引入模型的定义!!不能直接加载

If it is a self-defined model, save and load, you need to import the definition of the model !! 
cannot be loaded directly! 

上一章 16.初识Pytorch现有网络(existing network structures)的使用(use)与修改(modify)
下一章 18.初识Pytorch完整的模型训练套路-合在一个.py文件中 Complete model training routine - in one .py file

你可能感兴趣的:(手把手学习Pytorch,pytorch,python,深度学习,人工智能,神经网络)