pytorch模型的保存与加载

原创

已于 2023-12-21 23:07:29 修改 · 1w 阅读

87 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

于 2023-07-28 10:37:26 首次发布

PyTorch提供了三种保存和加载模型的方法：1)保存和加载模型参数（推荐），通常用于常规模型存储；2)保存和加载整个模型，包含模型结构和参数；3)导出和加载ONNX格式模型，用于跨框架和硬件平台的模型交换。使用torch.save和torch.load函数可以方便地进行模型的序列化和反序列化，而state_dict则用于存储模型的参数。ONNX格式允许模型在不同深度学习框架间转换。

1 pytorch保存和加载模型的三种方法

PyTorch提供了三种方式来保存和加载模型，在这三种方式中，加载模型的代码和保存模型的代码必须相匹配，才能保证模型的加载成功。通常情况下，使用第一种方式（保存和加载模型状态字典）更加常见，因为它更轻量且不依赖于特定的模型类。

1.1 仅保存和加载模型参数（推荐）

1.1.1 保存模型参数

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(128, 16), nn.ReLU(), nn.Linear(16, 1))

# 保存整个模型
torch.save(model.state_dict(), 'sample_model.pt')

1.1.2 加载模型参数

import torch
import torch.nn as nn

# 下载模型参数 并放到模型中
loaded_model = nn.Sequential(nn.Linear(128, 16), nn.ReLU(), nn.Linear(16, 1))
loaded_model.load_state_dict(torch.load('sample_model.pt'))
print(loaded_model)

显示如下：

Sequential(
  (0): Linear(in_features=128, out_features=16, bias=True)
  (1): ReLU()
  (2): Linear(in_features=16, out_features=1, bias=True)
)

state_dict：PyTorch中的state_dict是一个python字典对象，将每个层映射到其参数Tensor。state_dict对象存储模型的可学习参数，即权重和偏差，并且可以非常容易地序列化和保存。

1.2 保存和加载整个模型

1.2.1 保存整个模型

import torch
import torch.nn as nn

net = nn.Sequential(nn.Linear(128, 16), nn.ReLU(), nn.Linear(16, 1))

# 保存整个模型，包含模型结构和参数
torch.save(net, 'sample_model.pt')

1.2.2 加载整个模型

import torch
import torch.nn as nn

# 加载整个模型，包含模型结构和参数
loaded_model = torch.load('sample_model.pt')
print(loaded_model)

显示如下：

Sequential(
  (0): Linear(in_features=128, out_features=16, bias=True)
  (1): ReLU()
  (2

最低0.47元/天解锁文章