（多输入+动态维度）整理的自定义神经网络pt转onnx过程的python代码，记录了pt文件转onnx全过程，简单的修改即可应用。

pt文件转onnx步骤

  • 1、编写预处理代码
  • 2、用onnxruntime导出onnx
  • 3、对导出的模型进行检查
  • 4、推理onnx模型，查看输出是否一致
  • 5、对onnx模型的输出进行处理，显示cv图像
  • 6、编辑主函数进行测试

1、编写预处理代码

预处理代码 与torch模型的预处理代码一样

def preprocess(img):img = (cv2.cvtColor(img, cv2.COLOR_BGR2RGB)).transpose(2, 0, 1)img = np.expand_dims(img, 0)sh_im = img.shapeif sh_im[2]%2==1:img = np.concatenate((img, img[:, :, -1, :][:, :, np.newaxis, :]), axis=2)if sh_im[3]%2==1:img = np.concatenate((img, img[:, :, :, -1][:, :, :, np.newaxis]), axis=3)img = normalize(img)img = torch.Tensor(img)return img

2、用onnxruntime导出onnx

def export_onnx(net, model_path, img, nsigma, onnx_outPath):nsigma /= 255.if torch.cuda.is_available():state_dict = torch.load(model_path)model = net.cuda()dtype = torch.cuda.FloatTensorelse:state_dict = torch.load(model_path, map_location='cpu')state_dict = remove_dataparallel_wrapper(state_dict)model = netdtype = torch.FloatTensorimg = Variable(img.type(dtype))nsigma = Variable(torch.FloatTensor([nsigma]).type(dtype))# 我这里预训练权重中参数名字与网络名字不同# 相同的话可直接load_state_dict(state_dict)new_state_dict = {}for k, v in state_dict.items():new_state_dict[k[7:]] = vmodel.load_state_dict(new_state_dict)# 设置onnx的输入输出列表，多输入多输出就设置多个input_list = ['input', 'nsigma']output_list = ['output']# onnx模型导出# dynamic_axes为动态维度，如果自己的输入输出是维度变化的建议设置，否则只能输入固定维度的tensortorch.onnx.export(model, (img, nsigma), onnx_outPath, verbose=True, opset_version=11, export_params=True, input_names=input_list, output_names=output_list, dynamic_axes={'input_img': {0: 'batch', 1: 'channel', 2: 'height', 3: 'width'}, 'output': {0: 'batch', 1: 'channel', 2: 'height', 3: 'width'}})

导出结果

3、对导出的模型进行检查

此处为检查onnx模型节点，后面如果onnx算子不支持转engine时，方便定位节点，找到不支持的算子进行修改

def check_onnx(onnx_model_path):model = onnx.load(onnx_model_path)onnx.checker.check_model((model))print(onnx.helper.printable_graph(model.graph))

下面贴出输出结果


netron可视化

4、推理onnx模型，查看输出是否一致

def run_onnx(onnx_model_path, test_img, nsigma):nsigma /= 255.with torch.no_grad:# 这里默认是cuda推理torch.cuda.FloatTensorimg = Variable(test_img.type(torch.cuda.FloatTensor))nsigma = Variable(torch.FloatTensor([nsigma]).type(torch.cuda.FloatTensor))# 设置GPU推理device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")providers = ['CUDAExecutionProvider'] if device != "cpu" else ['CPUExecutionProvider']# 通过创建onnxruntime session来运行onnx模型ort_session = ort.InferenceSession(onnx_model_path, providers=providers)output = ort_session.run(output_names=['output'], input_feed={'input_img': np.array(img.cpu(), dtype=np.float32),'nsigma':np.array(nsigma.cpu(), dtype=np.float32)})return output

5、对onnx模型的输出进行处理，显示cv图像

def postprocess(img, img_noise_estime):out = torch.clamp(img-img_noise_estime, 0., 1.)outimg = variable_to_cv2_image(out)cv2.imshow(outimg)

6、编辑主函数进行测试

def main():################################onnx模型导出################################ pt权重路径：自己的路径 + mypt.ptmodel_path = "D:/python/ffdnet-pytorch/models/net_rgb.pth"# export onnx模型时输入进去数据，用于onnx记录网络的计算过程export_feed_path = "D:/python/ffdnet-pytorch/noisy.png"# onnx模型导出的路径onnx_outpath = "D:/python/ffdnet-pytorch/models/myonnx.onnx"# 实例化自己的网络模型并设置输入参数net = FFDNet(num_input_channels=3)nsigma = 25# onnx 导出img = cv2.imread(export_feed_path)input = preprocess(img)export_onnx(net, model_path, input, nsigma, onnx_outpath)print("export success!")################################检查onnx模型###############################check_onnx(onnx_outpath)# netron可视化网络，可视化用节点记录的网络推理流程netron.start(onnx_outpath)################################运行onnx模型################################ 此处过程是数据预处理 ---> 调用run_onnx函数 ---> 对模型输出后处理# 具体代码就不再重复了

#完整代码

import timeimport netronimport cv2import torchimport onnximport numpy as npfrom torch.autograd import Variableimport onnxruntime as ortfrom models import FFDNetfrom utils import remove_dataparallel_wrapper, normalize, variable_to_cv2_image# 此处为预处理代码 与torch模型的预处理代码一样def preprocess(img):img = (cv2.cvtColor(img, cv2.COLOR_BGR2RGB)).transpose(2, 0, 1)img = np.expand_dims(img, 0)sh_im = img.shapeif sh_im[2]%2==1:img = np.concatenate((img, img[:, :, -1, :][:, :, np.newaxis, :]), axis=2)if sh_im[3]%2==1:img = np.concatenate((img, img[:, :, :, -1][:, :, :, np.newaxis]), axis=3)img = normalize(img)img = torch.Tensor(img)return img# 此处为onnx模型导出的代码，包括torch模型的pt权重加载，onnx模型的导出def export_onnx(net, model_path, img, nsigma, onnx_outPath):nsigma /= 255.if torch.cuda.is_available():state_dict = torch.load(model_path)model = net.cuda()dtype = torch.cuda.FloatTensorelse:state_dict = torch.load(model_path, map_location='cpu')state_dict = remove_dataparallel_wrapper(state_dict)model = netdtype = torch.FloatTensorimg = Variable(img.type(dtype))nsigma = Variable(torch.FloatTensor([nsigma]).type(dtype))# 我这里预训练权重中参数名字与网络名字不同# 相同的话可直接load_state_dict(state_dict)new_state_dict = {}for k, v in state_dict.items():new_state_dict[k[7:]] = vmodel.load_state_dict(new_state_dict)# 设置onnx的输入输出列表，多输入多输出就设置多个input_list = ['input', 'nsigma']output_list = ['output']# onnx模型导出# dynamic_axes为动态维度，如果自己的输入输出是维度变化的建议设置，否则只能输入固定维度的tensortorch.onnx.export(model, (img, nsigma), onnx_outPath, verbose=True, opset_version=11, export_params=True,input_names=input_list, output_names=output_list,dynamic_axes={'input_img': {0: 'batch', 1: 'channel', 2: 'height', 3: 'width'},'output': {0: 'batch', 1: 'channel', 2: 'height', 3: 'width'}})# 此处为检查onnx模型节点，后面如果onnx算子不支持转engine时，方便定位节点，找到不支持的算子进行修改def check_onnx(onnx_model_path):model = onnx.load(onnx_model_path)onnx.checker.check_model((model))print(onnx.helper.printable_graph(model.graph))# 此处为推理onnx模型的代码，检查输出是否跟torch模型相同def run_onnx(onnx_model_path, test_img, nsigma):nsigma /= 255.with torch.no_grad:# 这里默认是cuda推理torch.cuda.FloatTensorimg = Variable(test_img.type(torch.cuda.FloatTensor))nsigma = Variable(torch.FloatTensor([nsigma]).type(torch.cuda.FloatTensor))# 设置GPU推理device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")providers = ['CUDAExecutionProvider'] if device != "cpu" else ['CPUExecutionProvider']# 通过创建onnxruntime session来运行onnx模型ort_session = ort.InferenceSession(onnx_model_path, providers=providers)output = ort_session.run(output_names=['output'], input_feed={'input_img': np.array(img.cpu(), dtype=np.float32), 'nsigma':np.array(nsigma.cpu(), dtype=np.float32)})return output# 此处是后处理代码，将onnx模型的输出处理成可显示cv图像# 与torch模型的后处理一样def postprocess(img, img_noise_estime):out = torch.clamp(img-img_noise_estime, 0., 1.)outimg = variable_to_cv2_image(out)cv2.imshow(outimg)def main():################################onnx模型导出################################ pt权重路径：自己的路径 + mypt.ptmodel_path = "D:/python/ffdnet-pytorch/models/net_rgb.pth"# export onnx模型时输入进去数据，用于onnx记录网络的计算过程export_feed_path = "D:/python/ffdnet-pytorch/noisy.png"# onnx模型导出的路径onnx_outpath = "D:/python/ffdnet-pytorch/models/myonnx.onnx"# 实例化自己的网络模型并设置输入参数net = FFDNet(num_input_channels=3)nsigma = 25# onnx 导出img = cv2.imread(export_feed_path)input = preprocess(img)export_onnx(net, model_path, input, nsigma, onnx_outpath)print("export success!")################################检查onnx模型###############################onnx_outpath = "D:/python/ffdnet-pytorch/models/myonnx.onnx"check_onnx(onnx_outpath)# netron可视化网络，可视化用节点记录的网络推理流程netron.start(onnx_outpath)################################运行onnx模型################################ 此处过程是数据预处理 ---> 调用run_onnx函数 ---> 对模型输出后处理# 具体代码就不再重复了if __name__ == '__main__':main()