pytorch教程(1.7)——保存和加载模型

最新推荐文章于 2025-03-18 19:36:33 发布

leetteel

最新推荐文章于 2025-03-18 19:36:33 发布

阅读量461

点赞数 1

分类专栏： # Pytorch 深度学习文章标签： pytorch python

版权

20 篇文章

订阅专栏

20 篇文章

订阅专栏

本文介绍了如何在PyTorch中保存和加载模型权重，确保在推理时设置模型为评估模式。此外，还展示了如何将模型导出到ONNX格式，以便在不同平台和语言上进行推理。关键步骤包括使用state_dict保存和加载权重，以及使用torch.onnx.export导出到ONNX。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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

import torch 
import torch.onnx as onnx 
import torchvision.models as models

保存和加载模型权重

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

一定要在推理之前调用 model.eval() 方法，以将 dropout 和批量归一化层设置为评估模式。不这样做会产生不一致的推理结果。

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

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

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

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

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

PyTorch 还具有本机 ONNX 导出支持。然而，鉴于 PyTorch 执行图的动态特性，导出过程必须遍历执行图以生成持久化的 ONNX 模型。出于这个原因，应该将适当大小的测试变量传递给导出例程（在我们的例子中，我们将创建一个正确大小的虚拟零张量）：

input_image = torch.zeros((1,3,224,224))
onnx.export(model, input_image, 'model.onnx')

您可以使用 ONNX 模型做很多事情，包括在不同平台和不同编程语言上运行推理。有关更多详细信息，我们建议访问 ONNX 教程。