保存和加载模型 ¶

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

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

import torch
import torchvision.models as models

保存和加载模型权重 ¶

PyTorch 模型将学习到的参数存储在一个内部状态字典中，称为 state_dict 。这些可以通过 torch.save 方法进行持久化：

model = models.vgg16(weights='IMAGENET1K_V1')
torch.save(model.state_dict(), 'model_weights.pth')

要加载模型权重，首先需要创建相同模型的实例，然后使用 load_state_dict() 方法加载参数。

在下面的代码中，我们将 weights_only=True 设置为限制在反序列化过程中执行的功能，仅限于加载权重所需的那些。在加载权重时使用 weights_only=True 被认为是最佳实践。

model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model
model.load_state_dict(torch.load('model_weights.pth', weights_only=True))
model.eval()

备注

在推理之前务必调用 model.eval() 方法，将 dropout 和批量归一化层设置为评估模式。未能这样做将导致推理结果不一致。

保存和加载具有形状的模型 ¶

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

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

然后我们可以像下面这样加载模型。

如在保存和加载 torch.nn.Modules 中所述，保存 state_dict 被认为是最佳实践。然而，下面我们使用 weights_only=False ，因为这涉及到加载模型，这是 torch.save 的遗留用例。

model = torch.load('model.pth', weights_only=False),

备注

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

保存和加载模型 ¶