一键转换onnx、动态batchsize、测试onnx、判断转换前后精度是否一致、将onnx模型的每个节点特征图信息输出（超详细合集、多输入）

原创

已于 2022-11-07 19:17:05 修改 · 4.5k 阅读

36 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #python

于 2022-11-07 19:15:29 首次发布

本文详细介绍如何将PyTorch模型转换为ONNX格式，并提供精度验证方法及特征图输出技巧。

最近项目用到onnx,对onnx内容进行总结，每一种都是测过可行，可以根据对应情况修改对应内容。

一、转换onnx

1. pt格式权重转onnx

#双输入，单输出
import torch.onnx
#根据模型图片输入设定input的h和w的尺寸大小，我是512
input = (torch.randn([1,3,512,512]).cuda(), torch.randn([1,3,512,512]).cuda())  
input_names = ['input_1', 'input_2']
output_names = ['output']

weight_path = './onnx/*****.pt'
model=torch.load(weight_path)
model = model.cuda()
model.eval()
print('load model over')

torch.onnx.export(model,
                  args=input,
                  f = './onnx/*****.onnx',
                  opset_version=11,
                  input_names=input_names,
                  output_names=output_names
                  )
print('to onnx success!')

2.pth格式权重转onnx

import torch.onnx
from model.siamunet_diff.SiamUNet_Diff import SiamUnet_diff #根据自己模型定义

####转换onnx####
input = (torch.randn([1,3,512,512]).cuda(), torch.randn([1,3,512,512]).cuda())

input_names = ['input_1', 'input_2']
output_names = ['output']

weight_path = '/home/*****.pth'
model = SiamUnet_diff(input_nbr=3, label_nbr=2).cuda()#根据自己模型定义

model.load_state_dict(torch.load(weight_path))
model.eval()#固定住dropout、 batchnorm

print('load model over')