PyTorch入门实践：PyTorch模型保存与加载入门

保存并加载模型

学习目标

在本次实验中，学员将了解在深度学习中，如何通过保存、加载和运行模型预测来保持模型状态，并将模型置于评估模式，以保证推理结果的确定性与一致性。

相关知识点

• PyTorch保存并加载模型

学习内容

1 PyTorch保存并加载模型

1.1 实验准备

#下载模型(pth文件)

!wget https://model-community-picture.obs.cn-north-4.myhuaweicloud.com/ascend-zone/notebook_models/7a98b34a258c11f0bd82fa163edcddae/vgg16-397923af.pth

%matplotlib inline
import torch
import torchvision.models as models

1.2 保存和加载模型权重

PyTorch 模型将学习的参数存储在称为 state_dict 的内部状态字典中。这些可以通过 torch.save 方法：

1.2 保存和加载模型权重

PyTorch 模型将学习的参数存储在称为 state_dict 的内部状态字典中。这些可以通过 torch.save 方法：

model = models.vgg16(weights=None)
state_dict = torch.load('./vgg16-397923af.pth', weights_only=True)
model.load_state_dict(state_dict)
torch.save(model.state_dict(), 'model_weights.pth')

要加载模型权重，需要先创建同一模型的实例，然后使用 load_state_dict（） 方法加载参数。
在下面的代码中，设置 weights_only=True 以将解封期间执行的函数限制为仅加载权重所需的函数。在加载权重时，使用 weights_only=True 被认为是最佳实践。

model2 = models.vgg16(weights=None)
model2.load_state_dict(torch.load('model_weights.pth', weights_only=True))
model2.eval()

注意:

请务必在推理之前调用 model.eval（） 方法，以将 dropout 和 batch normalization 层设置为评估模式。如果不这样做，将产生不一致的推理结果。

1.3 保存和加载带有 shape 的模型

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

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

然后，可以加载模型，如下所示。
保存 state_dict 被认为是最佳实践。但是，下面使用 weights_only=False，因为这涉及加载模型，这是 torch.save 的用例。

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

注意:

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