Pytorch | Pytorch格式 .pt .pth .bin .onnx 详解

Pytorch是深度学习领域中非常流行的框架之一，支持的模型保存格式包括.pt和.pth .bin .onnx。这几种格式的文件都可以保存Pytorch训练出的模型，但是它们的区别是什么呢？

模型的保存与加载到底在做什么？

我们在使用pytorch构建模型并且训练完成后，下一步要做的就是把这个模型放到实际场景中应用，或者是分享给其他人学习、研究、使用。因此，我们开始思考一个问题，提供哪些模型信息，能够让对方能够完全复现我们的模型？

• 模型代码：

• • （1）包含了我们如何定义模型的结构，包括模型有多少层/每层有多少神经元等等信息；
  • （2）包含了我们如何定义的训练过程，包括epoch batch_size等参数；
  • （3）包含了我们如何加载数据和使用；
  • （4）包含了我们如何测试评估模型。

• 模型参数：提供了模型代码之后，对方确实能够复现模型，但是运行的参数需要重新训练才能得到，而没有办法在我们的模型参数基础上继续训练，因此对方还希望我们能够把模型的参数也保存下来给对方。

• • （1）包含model.state_dict()，这是模型每一层可学习的节点的参数，比如weight/bias；
  • （2）包含optimizer.state_dict()，这是模型的优化器中的参数；
  • （3）包含我们其他参数信息，如epoch/batch_size/loss等。

• 数据集：

• • （1）包含了我们训练模型使用的所有数据；
  • （2）可以提示对方如何去准备同样格式的数据来训练模型。

• 使用文档：

• • （1）根据使用文档的步骤，每个人都可以重现模型；
  • （2）包含了模型的使用细节和我们相关参数的设置依据等信息。

可以看到，根据我们提供的模型代码/模型参数/数据集/使用文档，我们就可以有理由相信对方是有手就会了，那么目的就达到了。

现在我们反转一下思路，我们希望别人给我们提供模型的时候也能够提供这些信息，那么我们就可以拿捏住别人的模型了。

为什么要约定格式？

根据上一段的思路，我们知道模型重现的关键是模型结构/模型参数/数据集，那么我们提供或者希望别人提供这些信息，需要一个交流的规范，这样才不会1000个人给出1000种格式，而 .pt .pth .bin 以及 .onnx 就是约定的格式。

torch.save: Saves a serialized object to disk. This function uses Python’s pickle utility for serialization. Models, tensors, and dictionaries of all kinds of objects can be saved using this function.

不同的后缀只是用于提示我们文件可能包含的内容，但是具体的内容需要看模型提供者编写的README.md才知道。而在使用torch.load()方法加载模型信息的时候，并不是根据文件的后缀进行的读取，而是根据文件的实际内容自动识别的，因此对于torch.load()方法而言，不管你把后缀改成是什么，只要文件是对的都可以读取。

torch.load: Uses pickle’s unpickling facilities to deserialize pickled object files to memory. This function also facilitates the device to load the data into (see Saving & Loading Model Across Devices).

顺便提一下，“一切皆文件”的思维才是正确打开计算机世界的思维方式，文件后缀只作为提示作用，在Windows系统中也会用于提示系统默认如何打开或执行文件，除此之外，文件后缀不应该成为我们认识和了解文件阻碍。

格式汇总

下面是一个整理了 .pt、.pth、.bin、ONNX 和 TorchScript 等 PyTorch 模型文件格式的表格：

.pt .pth格式

一个完整的Pytorch模型文件，包含了如下参数：

• model_state_dict：模型参数
• optimizer_state_dict：优化器的状态
• epoch：当前的训练轮数
• loss：当前的损失值

下面是一个.pt文件的保存和加载示例（注意，后缀也可以是 .pth ）：

• .state_dict()：包含所有的参数和持久化缓存的字典，model和optimizer都有这个方法
• torch.save()：将所有的组件保存到文件中

模型保存

import torch
import torch.nn as nn

# 定义一个简单的模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
loss = 0.55

# 保存模型
state = {
    'epoch': 10,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    'loss': loss
}
checkpoint="./PATH/model-1.pt"
torch.save(state,checkpoint)

模型加载

import torch
import torch.nn as nn

# 定义同样的模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

# 加载模型
model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()

.bin格式

.bin文件是一个二进制文件，可以保存Pytorch模型的参数和持久化缓存。.bin文件的大小较小，加载速度较快，因此在生产环境中使用较多。

下面是一个.bin文件的保存和加载示例（注意：也可以使用 .pt .pth 后缀）：

保存模型

import torch
import torch.nn as nn

# 定义一个简单的模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

model = Net()
# 保存参数到.bin文件
checkpoint="./PATH/model.bin"
torch.save(model.state_dict(), checkpoint)

加载模型

import torch
import torch.nn as nn

# 定义相同的模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

# 加载.bin文件
model = Net()
checkpoint="./PATH/model.bin"
model.load_state_dict(torch.load(checkpoint))
model.eval()

.onnx格式

上述保存的文件可以通过PyTorch提供的torch.onnx.export函数转化为ONNX格式，这样可以在其他深度学习框架中使用PyTorch训练的模型。转化方法如下：

import torch
import torch.onnx

# 将模型保存为.bin文件
model = torch.nn.Linear(3, 1)
torch.save(model.state_dict(), "model.bin")
# torch.save(model.state_dict(), "model.pt")
# torch.save(model.state_dict(), "model.pth")

# 将.bin文件转化为ONNX格式
model = torch.nn.Linear(3, 1)
model.load_state_dict(torch.load("model.bin"))
# model.load_state_dict(torch.load("model.pt"))
# model.load_state_dict(torch.load("model.pth"))
example_input = torch.randn(1, 3)
torch.onnx.export(model, example_input, "model.onnx", input_names=["input"], output_names=["output"])

加载ONNX格式的代码可以参考以下示例代码：

import onnx
import onnxruntime

# 加载ONNX文件
onnx_model = onnx.load("model.onnx")

# 将ONNX文件转化为ORT格式
ort_session = onnxruntime.InferenceSession("model.onnx")

# 输入数据
input_data = np.random.random(size=(1, 3)).astype(np.float32)

# 运行模型
outputs = ort_session.run(None, {"input": input_data})

# 输出结果
print(outputs)

注意，需要安装onnx和onnxruntime两个Python包。此外，还需要使用numpy等其他常用的科学计算库。

直接保存完整模型

可以看出来，我们在之前的报错方式中，都是保存了.state_dict()，但是没有保存模型的结构，在其他地方使用的时候，必须先重新定义相同结构的模型（或兼容模型），才能够加载模型参数进行使用，如果我们想直接把整个模型都保存下来，避免重新定义模型，可以按如下操作：

# 保存模型
PATH = "entire_model.pt"
# PATH = "entire_model.pth"
# PATH = "entire_model.bin"
torch.save(model, PATH)

加载模型

# 加载模型
model = torch.load("entire_model.pt")
model.eval()

结语

本文介绍了pytorch可以导出的模型的几种后缀格式，但是模型导出的关键并不是后缀，而是到处时候提供的信息到底是什么，只要知道了模型的model.state_dict()和optimizer.state_dict()，以及相应的epoch batch_size loss等信息，我们就能够重建出模型，至于要导出哪些信息，就取决于你了，务必在readme.md中写清楚，你导出了哪些信息。

本栏目推荐文章
• 手把手教你MongoDB 数据库连接URL 格式、authSource参数
• 字符串格式化
• [JMeter] JMeter的测试报告格式转换(.jtl => html)
• 2024年更新「GIS数据」全国的GeoJSON、shp格式数据下载获取(精确到乡镇街道级)
• 第六天：shell脚本编程基础格式规范及变量
• axios中使用qs.stringify格式化get请求参数
• 使用daplink下载任意格式固件
• ARM ELF文件格式
• stm32 文件格式对比
• ELF文件格式

Pytorch 格式 onnx bin pthpytorch格式onnx bin 卷积pytorch问题onnx 深度pytorch onnx 模型pytorch unet onnx 模型pytorch onnx 格式 文件bin axf pth 模型 文件plk pth onnx onnx-mlir