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

原创已于 2022-01-29 00:05:16 修改 · 2.7k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #python #深度学习 #人工智能 #神经网络

于 2022-01-21 19:26:24 首次发布

手把手学习Pytorch 专栏收录该内容

23 篇文章

订阅专栏

本文介绍了两种Pytorch模型的保存方法：一是保存整个模型；二是仅保存模型参数。同时展示了如何加载这两种不同形式的保存文件，并强调了自定义模型在保存与加载时需要注意的问题。

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

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：
在这里插入图片描述

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

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:
在这里插入图片描述 …

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

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:
在这里插入图片描述

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

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