PyTorchTutorials【pytorch官方教程中英文详解】8SaveandLoadModel

在文章PyTorch-Tutorials【pytorch官方教程中英文详解】- 7 Optimization中介绍了模型如何优化&＃xff0c;包括构造损失函数和优化器等&＃xff0c;接下来看看如何保存已经优化的模型&＃xff0c;以及如何载入保存的模型。

原文链接&＃xff1a;Save and Load the Model — PyTorch Tutorials 1.10.1&＃43;cu102 documentation

SAVE AND LOAD THE MODEL

In this section we will look at how to persist model state with saving, loading and running model predictions.

【在本节中&＃xff0c;我们将看看如何通过保存、加载和运行模型预测来持久化模型状态。】

import torch
import torchvision.models as models
1 Saving and Loading Model Weights

PyTorch models store the learned parameters in an internal state dictionary, called state_dict. These can be persisted via the torch.save method:

【PyTorch模型将学习到的参数存储在一个名为state_dict的内部状态字典中。这些可以通过torch.save方法:】

model &＃61; models.vgg16(pretrained&＃61;True)
torch.save(model.state_dict(), &＃39;model_weights.pth&＃39;)

To load model weights, you need to create an instance of the same model first, and then load the parameters using load_state_dict() method.

【要加载模型权重&＃xff0c;您需要首先创建同一个模型的实例&＃xff0c;然后使用load_state_dict()方法加载参数。】

model &＃61; models.vgg16() # we do not specify pretrained&＃61;True, i.e. do not load default weights
model.load_state_dict(torch.load(&＃39;model_weights.pth&＃39;))
model.eval()

NOTE

• be sure to call model.eval() method before inferencing to set the dropout and batch normalization layers to evaluation mode. Failing to do this will yield inconsistent inference results.
• 【在推断将退出和批处理规范化层设置为求值模式之前&＃xff0c;请确保调用model.eval()方法。如果不这样做&＃xff0c;将产生不一致的推理结果。】

When loading model weights, we needed to instantiate the model class first, because the class defines the structure of a network. We might want to save the structure of this class together with the model, in which case we can pass model (and not model.state_dict()) to the saving function:

【当加载模型权重时&＃xff0c;我们需要首先实例化模型类&＃xff0c;因为类定义了网络的结构。我们可能希望将这个类的结构与模型一起保存&＃xff0c;在这种情况下&＃xff0c;我们可以将model(而不是model.state_dict())传递给保存函数:】

torch.save(model, &＃39;model.pth&＃39;)

We can then load the model like this:

【然后我们可以像这样加载模型:】

model &＃61; torch.load(&＃39;model.pth&＃39;)

NOTE

• This approach uses Python pickle module when serializing the model, thus it relies on the actual class definition to be available when loading the model.
• 【这种方法在序列化模型时使用Python pickle模块&＃xff0c;因此它依赖于在加载模型时可用的实际类定义。】

Saving and Loading a General Checkpoint in PyTorch

恭喜&＃xff01;整个【pytorch官方教程中英文详解】系列更新完毕&＃xff0c;一共八节。

说明&＃xff1a;记录学习笔记&＃xff0c;如果错误欢迎指正&＃xff01;写文章不易&＃xff0c;转载请联系我。