pytorch模型转onnx模型,加速推理（pytorch模型(.pt)转onnx模型(.onnx)的方法详解(1)）

1. pytorch模型转换到onnx模型

2.运行onnx模型

3.比对onnx模型和pytorch模型的输出结果

 我这里重点是第一点和第二点
  
  
  
  
  
  
，第三部分  比较容易

首先你要安装 依赖库：onnx 和 onnxruntime 
  
  
  
  
  
 ，

pip install onnx pip install onnxruntime 进行安装

也可以使用清华源镜像文件安装  速度会快些            。

开始：

1. pytorch模型转换到onnx模型

pytorch 转 onnx 仅仅需要一个函数 torch.onnx.export 

torch.onnx.export(model, args, path, export_params, verbose, input_names, output_names, do_constant_folding, dynamic_axes, opset_version)

参数说明：

model——需要导出的pytorch模型 args——模型的输入参数 


 


 


 


 


 


 ，满足输入层的shape正确即可  
 
  
 
  
 
  
 
  
 
  
 
 。 path——输出的onnx模型的位置












。例如‘yolov5.onnx’
  
  
  
  
  
  。 export_params——输出模型是否可训练  
  
  
  
  
  
 。default=True
 
 
 
 
 
 
，表示导出trained model  
   
   
   
   
   
 ，否则untrained
 


 


 


 


 


 


。 verbose——是否打印模型转换信息
 
 
 
 
 
 。default=False   
   
   
   
   
   
 。 input_names——输入节点名称
 

 

 

 

 

 

。default=None

 

 

 

 

 

 。 output_names——输出节点名称             。default=None
 
 

 
 

 
 

 
 

 
 

 
 

。 do_constant_folding——是否使用常量折叠 

 

 

 

 

 

 ，默认即可











。default=True                   。 dynamic_axes——模型的输入输出有时是可变的

 

 

 

 

 

 
，如Rnn            ，或者输出图像的batch可变 
 

 
 

 
 

 
 

 
 

 
 

 ，可通过该参数设置 
 
 
 
 
 
 
 
 
 
 
 
。如输入层的shape为（b












，3                  ，h 
 
 
 
 
 
 
 
 
 
 
 
 ，w）


















，batch            ，height 
 
  
 
  
 
  
 
  
 
  
 
  ，width是可变的











，但是chancel是固定三通道


















。

格式如下 ：

1)仅list(int) dynamic_axes={‘input’:[0,2,3],‘output’:[0,1]}

2)仅dict dynamic_axes={‘input’:{0:‘batch’,2:‘height’,3:‘width’},‘output’:{0:‘batch’,1:‘c’}}

3)mixed dynamic_axes={‘input’:{0:‘batch’,2:‘height’,3:‘width’},‘output’:[0,1]} opset_version——opset的版本
  
  
  
  
  
  
，低版本不支持upsample等操作            。

转化代码：参考1：

import torch import torch.nn import onnx model = torch.load(best.pt) model.eval() input_names = [input] output_names = [output] x = torch.randn(1,3,32,32,requires_grad=True) torch.onnx.export(model, x, best.onnx, input_names=input_names, output_names=output_names, verbose=True)

 参考2：PlainC3AENetCBAM 是网络模型 
  
  
  
  
  
 ，如果你没有自己的网络模型 


 


 


 


 


 


 ，可能成功不了

import io import torch import torch.onnx from models.C3AEModel import PlainC3AENetCBAM device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def test(): model = PlainC3AENetCBAM() pthfile = r/home/joy/Projects/models/emotion/PlainC3AENet.pth loaded_model = torch.load(pthfile, map_location=cpu) # try: # loaded_model.eval() # except AttributeError as error: # print(error) model.load_state_dict(loaded_model[state_dict]) # model = model.to(device) #data type nchw dummy_input1 = torch.randn(1, 3, 64, 64) # dummy_input2 = torch.randn(1, 3, 64, 64) # dummy_input3 = torch.randn(1, 3, 64, 64) input_names = [ "actual_input_1"] output_names = [ "output1" ] # torch.onnx.export(model, (dummy_input1, dummy_input2, dummy_input3), "C3AE.onnx", verbose=True, input_names=input_names, output_names=output_names) torch.onnx.export(model, dummy_input1, "C3AE_emotion.onnx", verbose=True, input_names=input_names, output_names=output_names) if __name__ == "__main__": test()

直接将PlainC3AENetCBAM替换成需要转换的模型
 
 
 
 
 
 
，然后修改pthfile  
   
   
   
   
   
 ，输入和onnx模型名字然后执行即可  
 
  
 
  
 
  
 
  
 
  
 
 。

注意：上面代码中注释的dummy_input2 

 

 

 

 

 

 ，dummy_input3

 

 

 

 

 

 
，torch.onnx.export对应的是多个输入的例子












。

在转换过程中遇到的问题汇总

RuntimeError: Failed to export an ONNX attribute, since its not constant, please try to make things (e.g., kernel size) static if possible

在转换过程中遇到RuntimeError: Failed to export an ONNX attribute, since its not constant, please try to make things (e.g., kernel size) static if possible的错误
  
  
  
  
  
  。

我成功的案例            ，我直接把我训练的网络贴上 
 

 
 

 
 

 
 

 
 

 
 

 ，成功转换












，没有from **   import 模型名词这么委婉                  ，合法 
 
 
 
 
 
 
 
 
 
 
 
 ，我的比较粗暴

import torch import torch.nn import onnx from torchvision import transforms import torch.nn as nn from torch.nn import Sequential # 添加模型 # 设置数据转换方式 preprocess_transform = transforms.Compose([ transforms.ToTensor(), # 把数据转换为张量（Tensor） transforms.Normalize( # 标准化


















，即使数据服从期望值为 0            ，标准差为 1 的正态分布 mean=[0.5, ], # 期望 std=[0.5, ] # 标准差 ) ]) class CNN(nn.Module): # 从父类 nn.Module 继承 def __init__(self): # 相当于 C++ 的构造函数 # super() 函数是用于调用父类(超类)的一个方法 
 
  
 
  
 
  
 
  
 
  
 
  ，是用来解决多重继承问题的 super(CNN, self).__init__() # 第一层卷积层  
  
  
  
  
  
 。Sequential(意为序列) 括号内表示要进行的操作 self.conv1 = Sequential( nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) # 第二卷积层 self.conv2 = Sequential( nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) # 全连接层（Dense











，密集连接层） self.dense = Sequential( nn.Linear(7 * 7 * 128, 1024), nn.ReLU(), nn.Dropout(p=0.5), nn.Linear(1024, 10) ) def forward(self, x): # 正向传播 x1 = self.conv1(x) x2 = self.conv2(x1) x = x2.view(-1, 7 * 7 * 128) x = self.dense(x) return x # 训练 # 训练和参数优化 # 定义求导函数 def get_Variable(x): x = torch.autograd.Variable(x) # Pytorch 的自动求导 # 判断是否有可用的 GPU return x.cuda() if torch.cuda.is_available() else x # 判断是否GPU device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # device1 = torch.device(cpu) # 定义网络 model = CNN() loaded_model = torch.load(save_model/model.pth, map_location=cuda:0) model.load_state_dict(loaded_model) model.eval() input_names = [input] output_names = [output] # x = torch.randn(1,3,32,32,requires_grad=True) x = torch.randn(1, 1, 28, 28, requires_grad=True) # 这个要与你的训练模型网络输入一致
 


 


 


 


 


 


。我的是黑白图像 torch.onnx.export(model, x, save_model/model.onnx, input_names=input_names, output_names=output_names, verbose=True)

前提是你要准备好*.pth模型保持文件

输出结果：

graph(%input : Float(1, 1, 28, 28, strides=[784, 784, 28, 1], requires_grad=1, device=cpu), %dense.0.weight : Float(1024, 6272, strides=[6272, 1], requires_grad=1, device=cpu), %dense.0.bias : Float(1024, strides=[1], requires_grad=1, device=cpu), %dense.3.weight : Float(10, 1024, strides=[1024, 1], requires_grad=1, device=cpu), %dense.3.bias : Float(10, strides=[1], requires_grad=1, device=cpu), %33 : Float(64, 1, 3, 3, strides=[9, 9, 3, 1], requires_grad=0, device=cpu), %34 : Float(64, strides=[1], requires_grad=0, device=cpu), %36 : Float(128, 64, 3, 3, strides=[576, 9, 3, 1], requires_grad=0, device=cpu), %37 : Float(128, strides=[1], requires_grad=0, device=cpu)): %input.4 : Float(1, 64, 28, 28, strides=[50176, 784, 28, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%input, %33, %34) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\conv.py:443:0 %21 : Float(1, 64, 28, 28, strides=[50176, 784, 28, 1], requires_grad=1, device=cpu) = onnx::Relu(%input.4) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:1442:0 %input.8 : Float(1, 64, 14, 14, strides=[12544, 196, 14, 1], requires_grad=1, device=cpu) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%21) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:797:0 %input.16 : Float(1, 128, 14, 14, strides=[25088, 196, 14, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%input.8, %36, %37) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\conv.py:443:0 %25 : Float(1, 128, 14, 14, strides=[25088, 196, 14, 1], requires_grad=1, device=cpu) = onnx::Relu(%input.16) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:1442:0 %26 : Float(1, 128, 7, 7, strides=[6272, 49, 7, 1], requires_grad=1, device=cpu) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%25) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:797:0 %27 : Long(2, strides=[1], device=cpu) = onnx::Constant[value= -1 6272 [ CPULongType{2} ]]() # E:/paddle_project/Pytorch_Imag_Classify/zifu_fenlei/CNN/pt模型转onnx模型.py:51:0 %28 : Float(1, 6272, strides=[6272, 1], requires_grad=1, device=cpu) = onnx::Reshape(%26, %27) # E:/paddle_project/Pytorch_Imag_Classify/zifu_fenlei/CNN/pt模型转onnx模型.py:51:0 %input.20 : Float(1, 1024, strides=[1024, 1], requires_grad=1, device=cpu) = onnx::Gemm[alpha=1., beta=1., transB=1](%28, %dense.0.weight, %dense.0.bias) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\linear.py:103:0 %input.24 : Float(1, 1024, strides=[1024, 1], requires_grad=1, device=cpu) = onnx::Relu(%input.20) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:1442:0 %output : Float(1, 10, strides=[10, 1], requires_grad=1, device=cpu) = onnx::Gemm[alpha=1., beta=1., transB=1](%input.24, %dense.3.weight, %dense.3.bias) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\linear.py:103:0 return (%output)

输出结果的device  是CPU
  
  
  
  
  
  
，模型加载的时候是GPU
 
 
 
 
 
 。这就是转换的意义吧

2.运行onnx模型

import onnx import onnxruntime as ort model = onnx.load(best.onnx) onnx.checker.check_model(model) session = ort.InferenceSession(best.onnx) x=np.random.randn(1,3,32,32).astype(np.float32) # 注意输入type一定要np.float32!!!!! # x= torch.randn(batch_size,chancel,h,w) outputs = session.run(None,input = { input : x })

参考：

Pytorch模型转onnx模型实例_python_脚本之家 (jb51.net)

pytorch模型转onnx模型的方法详解_python_脚本之家 (jb51.net)