一键转换onnx、动态batchsize、测试onnx、判断转换前后精度是否一致、将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')

torch.onnx.export(model,
                  args=input,
                  f = '/home/CropLand-CD/onnx/simen_unet.onnx',
                  opset_version=11,
                  input_names=input_names,
                  output_names=output_names)
print('to onnx success!')

3. 设置动态batchsize转onnx

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

input = (torch.randn([1,3,512,512]).cuda(), torch.randn([1,3,512,512]).cuda())

#多输入多输出,若设置单输入，删减names内容即可
input_names = ['input_1', 'input_2']
output_names = ['output_1','output_2','output_3']

import torch.nn as nn
weight_path = '/home/******.pth'
model = SiamUnet_diff(input_nbr=3, label_nbr=2).cuda()

model.load_state_dict(torch.load(weight_path))
# model = nn.DataParallel(model).cuda() #如果多显卡训练需要加这个
model.eval()

print('load model over')

#可以定义动态宽高
# dynamic_axes = {'input_1': {0: 'batch_size', 1: 'channel', 2: 'height', 3: 'width'},
#                 'input_2': {0: 'batch_size', 1: 'channel', 2: 'height', 3: 'width'},
#                 'output_1': {0: 'batch_size', 1: 'channel', 2: 'height', 3: 'width'},
#                 'output_2': {0: 'batch_size', 1: 'channel', 2: 'height', 3: 'width'},
#                 'output_3': {0: 'batch_size', 1: 'channel', 2: 'height', 3: 'width'}}
#

dynamic_axes = {'input_1': {0: 'batch_size'},
                 'input_2': {0: 'batch_size'},
                 'output_1': {0: 'batch_size'},
                 'output_2': {0: 'batch_size'},
                 'output_3': {0: 'batch_size'}}

torch.onnx.export(model,
                  args=input,
                  f = '/home/CropLand-CD/onnx/simen_unet.onnx',
                  opset_version=11,
                  input_names=input_names,
                  output_names=output_names，
                  dynamic_axes = dynamic_axes)
print('to onnx success!')

二、测试onnx

###双输入图片测试
import onnxruntime as ort
import cv2
import numpy as np
from PIL import Image

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
def resize_img_2(path, size=(512,512)):
    img = Image.open(path)
    img = img.resize(size)
    img = np.asarray(img, dtype='float32')
    img = np.expand_dims(img.transpose(2, 0, 1), axis=0)
    # print('img', img.shape)
    return img

# 读取图片
path1 = './**/1.png'
path2 = './**/2.png'

img1 = resize_img_2(path1)
img2 = resize_img_2(path2)

# 加载 onnx
from onnxruntime.datasets import get_example
onnx_model = get_example(r'/home/*****.onnx')#一定要绝对路径
sess = ort.InferenceSession(onnx_model,providers=['CUDAExecutionProvider']) # 'CPUExecutionProvider'
input_name_1 = sess.get_inputs()[0].name
input_name_2 = sess.get_inputs()[1].name

outputs = sess.run(None, {input_name_1:img1, input_name_2:img2})

file_path = './result/'
cv2.imwrite(file_path + '******.png', outputs) #保存图片
print('success!')

三、判断转换前后精度是否一致

1. 测试模型转换是否正常

#测试转换模型是否有问题#
import onnx
onnx_path = './onnx/******.onnx'
onnx_model = onnx.load(onnx_path)
check = onnx.checker.check_model(onnx_model)
print('check', check)

check 为 none 则正常

2. 测试精度是否一致

import onnxruntime
import numpy as np
from onnxruntime.datasets import get_example
import torch
from model.CDNET import CDNET #根据自己模型定义

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

model = CDNet(n_class=2,img_size=512).cuda()
checkpoint = torch.load('/home/CropLand-CD/onnx/*****.pth')
# model = nn.DataParallel(model).cuda() #如果用了多张显卡训练
model.load_state_dict(checkpoint)

# 得到torch模型的输出
dummy_input = torch.randn(1, 3, 512, 512)
dummy_input_2 = torch.randn(1, 3, 512, 512)
model.eval()
with torch.no_grad():
    torch_out = model(dummy_input,dummy_input_2)

# 得到onnx模型的输出
onnx_model = get_example(r'/home/*****.onnx')  # 要写绝对路径
sess = onnxruntime.InferenceSession(onnx_model, providers=['CUDAExecutionProvider'])
# sess = ort.InferenceSession(onnx_path,providers=['CUDAExecutionProvider']) # 'CPUExecutionProvider'
input_name_1 = sess.get_inputs()[0].name
input_name_2 = sess.get_inputs()[1].name


onnx_out = sess.run(None, {input_name_1:to_numpy(dummy_input), input_name_2:to_numpy(dummy_input_2)})

# 判断输出结果是否一致，小数点后3位一致即可
print(np.testing.assert_almost_equal(to_numpy(torch_out), onnx_out, decimal=3))

四、将onnx模型的每个节点特征图信息输出

from model.network import CDNet
from torchsummary import summary
import torch
net = CDNet().cuda()
summary(net, input_size = [(3,512,512),(3,512,512)])
####两个输入，单输入对应删减####