Detailed explanation of the method of converting pytorch model (.pt) to onnx model (.onnx) (1)

1. Convert pytorch model to onnx model

2. Run the onnx model

3. Compare the output results of the onnx model and the pytorch model

My focus here is on the first and second points, the third part is easier

First you need to install dependent libraries: onnx and onnxruntime,

pip install onnx
pip install onnxruntime 进行安装

You can also use the Tsinghua source image file to install it faster.

start:

1. Convert pytorch model to onnx model

pytorch to onnx only needs a function torch.onnx.export

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

Parameter Description:

model - the pytorch model that needs to be exported
args——The input parameters of the model, the shape of the input layer is correct.
path - the location of the output onnx model. For example 'yolov5.onnx'.
export_params - whether the output model is trainable. default=True, means to export the trained model, otherwise untrained.
verbose - whether to print model conversion information. default=False.
input_names - Input node names. default=None.
output_names - Output node names. default=None.
do_constant_folding——Whether to use constant folding, the default is fine. default=True.
dynamic_axes——The input and output of the model are sometimes variable, such as Rnn, or the batch of the output image is variable, which can be set through this parameter. For example, the shape of the input layer is (b, 3, h, w), batch, height, and width are variable, but chance is a fixed three-channel.
The format is as follows:
1) only list(int) dynamic_axes={'input':[0,2,3],'output':[0,1]}
2) only dict<int, string> 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 version, lower versions do not support operations such as upsample.

Conversion code: Reference 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')

Reference 2: PlainC3AENetCBAM is a network model, if you do not have your own network model, you may not succeed

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()

Directly replace PlainC3AENetCBAM with the model that needs to be converted, then modify pthfile, input and onnx model name and execute.

Note: The dummy_input2, dummy_input3, and torch.onnx.export commented in the above code correspond to examples of multiple inputs.

Summary of problems encountered during conversion

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

During the conversion process, I encountered a RuntimeError: Failed to export an ONNX attribute, since it's not constant, please try to make things (eg, kernel size) static if possible error.

In my successful case, I directly pasted the network I trained and converted it successfully. It is not as euphemistic and legal as the from ** import model noun, but mine is more rude

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')

The premise is that you have to prepare the *.pth model to keep the file

Output result:

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)

The device that outputs the results is the CPU, and the device when the model is loaded is the GPU. That's what conversion means.

2. Run the onnx model

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 })

reference:

Pytorch model to onnx model example Develop Paper

Detailed explanation of the method of converting pytorch model to onnx model Develop Paper

Detailed explanation of the method of converting pytorch model (.pt) to onnx model (.onnx) (1)

Guess you like