pth权重文件转换为onnx权重文件

阿月浑子
学习笔记

发布于：2022年10月31日

字数：407字

时长：1分钟

参考博文：

1.pth文件转为onnx格式业精于勤。荒于嬉。的博客-CSDN博客.pth转onnx

2.Pytorch——初探onnx（1）解决upsamplebilinear2d转换问题零壹博弈的博客-CSDN博客

3.https://blog.csdn.net/magic_show_time/article/details/122476306

由于想利用zetane这个软件可视化目标检测模型里面的特征变化过程，zetane只支持onnx、h5、ZTN文件,而我的目标检测模型得到的权重文件是pth文件,因此需要用到权重转换

经过一系列的摸索，得出以下代码（亲测可用）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import torch.onnx
import onnxruntime as ort
from nets import ssd

# 创建.pth模型

model = ssd.SSD300(21,'vgg')
# 加载权重
model_path = './model_data/ep250-loss1.474-val_loss2.826.pth'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_statedict = torch.load(model_path, map_location=device)
model.load_state_dict(model_statedict)

model.to(device)
model.eval()

input_data = torch.randn(1, 3, 300, 300, device=device)

# 转化为onnx模型
input_names = ['input']
output_names = ['output']

torch.onnx.export(model, input_data, 'ssdv7.onnx', opset_version=11, verbose=True, input_names=input_names,
                  output_names=output_names)

遇到的问题：

如果你是按照我给的两篇文章写代码，你会发现很简单

1.模型无法加载，load_state_dict出现问题

1
2
3
4
5
# 记住，模型参数一定要写全，根据自己的网络模型写即可
# 错误代码
model = ssd.SSD300()
# 正确代码
model = ssd.SSD300(21, 'vgg')

2.权重文件因为模型中上采样导出失败问题

image-20221031111722457

按方法2解决即可，pytorch版本问题

1
2
3
4
5
6
7
# 错误代码
torch.onnx.export(model, input_data, 'ssdv7.onnx', opset_version=9, verbose=True, input_names=input_names, output_names=output_names)

# 正确代码
torch.onnx.export(model, input_data, 'ssdv7.onnx', opset_version=11, verbose=True, input_names=input_names,
                  output_names=output_names)

ok,接下来就是zetane加载模型了，看看最终的模型吧。

image-20221031112309197

更新于：2022年10月31日

文件转换