【问题解决】将pytorch模型及训练好的权重转换为onnx文件

目录

1 ONNX是什么

2 pytorch模型转换为onnx文件

3 可能遇到的问题

1 ONNX是什么

ONNX（Open Neural Network Exchange）是一个开放的、跨平台的深度学习模型表示和转换工具。它的目标是使不同深度学习框架之间能够无缝地交换和使用模型。ONNX 的设计思想是将深度学习模型表示为计算图，其中节点表示操作（如卷积、池化等），边表示数据流动。这个计算图可以在不同的深度学习框架之间共享和执行。

在跨平台和跨框架使用模型、模型加速、模型在客户处部署时，经常需要用到onnx格式的模型。

2 pytorch模型转换为onnx文件

这里提供一种自定义pytorch模型和训练参数打包转换为onnx文件的方法。

import torch.onnx
import onnx

from models import MyModel  # 导入你的模型

Load_Path = 'my_dict_path.tar'  # 训练好的模型权重文件保存位置
Save_Path = 'my_file_path.onnx'  # 想要存放onnx文件的位置

# 实例化模型并加载权重
model = MyModel()
model.load_state_dict(torch.load(Load_Path, map_location=torch.device('cuda')))
model.eval()

# 定义示例输入
input_sample = torch.rand(1, 3, 64, 224, 224)

# 将模型转换为 ONNX 格式
torch.onnx.export(model, input_sample, Save_Path, 
                  verbose=True, input_names=['input'], output_names=['output'])

print("Model converted to ONNX format and saved to", Save_Path)

定义示例输入应与实际输入大小匹配。 

verbose参数用于控制程序输出的详细程度。当设置为True时，程序会输出更多的信息和细节。当设置为False时，程序的输出会更加简洁，只显示必要的信息。

3 可能遇到的问题

在写的过程中遇到了一些问题。

1

仅直接导入torch.onnx运行可能会报错

File"C:\Users\***\AppData\Local\Programs\Python\Python39\lib\torch\onnx\_internal\onnx_proto_utils.py", line 221, in _add_onnxscript_fn raise errors.OnnxExporterError("Module onnx is not installed!") from e torch.onnx.errors.OnnxExporterError: Module onnx is not installed!

先使用pip install onnx安装onnx包即可。

2

RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device

这个错误通常发生在你尝试加载一个在不同 GPU 设备上训练过的模型参数时。由于当前系统只有一个 GPU 设备，而模型参数保存在了另一个 GPU 设备上，所以导致了这个错误。

为了解决这个问题，可以在使用torch.load函数时，通过map_location参数将模型参数映射到当前可用的设备。