Pytorch官方文档学习笔记 —— 6. Save & Load Model

目录

1. 保存和加载模型

2. 保存和加载模型权重

3. 保存和加载带形状的模型

4. 相关教程

---

1. 保存和加载模型

在本节中，我们将了解如何通过保存、加载和运行模型预测来保持模型状态。

import torch
import torchvision.models as models

2. 保存和加载模型权重

        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()

Note：请确保在推理前调用 model.eval() 方法，将 dropout 和 batch normalization 层设置为评估模式。如果不这样做，将产生不一致的推理结果。

3. 保存和加载带形状的模型

        加载模型权重时，我们需要先实例化模型类，因为该类定义了网络的结构。我们可能希望将此类的结构与模型一起保存，在这种情况下，我们可以将模型（而不是 model.state_dict()）传递给保存函数：

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

然后我们可以像这样加载模型：

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

Note：这种方法在序列化模型时使用 Python pickle 模块，因此它依赖于在加载模型时可用的实际类定义。

4. 相关教程

        在 PyTorch 中保存和加载常规检查点。