ONNX model conversion and onnxruntime reasoning

ONNX is used as an intermediate variable for model deployment. Both tensorRT and tfsavemodel need to convert the pytorch or tensorflow model to ONNX first. ONNX can also be used as an intermediate variable for reasoning through onnxruntime.

ONNX model conversion

import torch
import torch.onnx
import onnx
weight=''
model = net()
model.load_state_dict(torch.load(weight, map_location='cpu'))
model.eval()
dummy_input = torch.randn(1, 3, 640, 640) # 需要预定义一个样例张量，张量的尺寸和需要输入到网络的尺寸保持一致
output_name='1.onnx'
torch.onnx.export(
        model, # 模型
        dummy_input, # 输入数据张量
        output_name, # onnx模型地址
        input_names=["images"], # 指定输入张量的name
        output_names=["output"], # 指定输出张量的name，如果是
        dynamic_axes={'images': {0: 'batch_size', 1: 'channel', 2: "height", 3: 'width'},'output': {0: 'batch'}},# 选择输入张量可以是动态尺寸
        opset_version= 11 # onnx的version
    )

• weight should be a pre-trained model file ending in .pth or pt
• The acquisition of the model needs to be changed, that is, the instantiation of your model class can be changed by itself
• Since the weights are saved in different ways during training, there are two types. One is to keep only the weights, then you need to use load_state_dict to load the read weight file to the model. The other is to save the weight and the model together, so that the process of obtaining the model is also omitted. You can directly model=torch.load(weight) to directly obtain the model with the loaded weight.
• model.eval() must call eval when doing inference.
• When dummy_input converts onnx, you must first set an input shape tensor. If the dynamic_axes parameter below is set to None, then this tensor must be exactly the same as the shape accepted by the inference model.
• output_name needs to specify the path and name of the onnx file, you can define it yourself, for example /user/onnx/model.onnx
• Torch.onnx.export() converts torch to onnx. In addition to some parameters mentioned above, it also needs to set an input_names and output_names
• Dynamic input can be set during model conversion, that is, your input can be different every time, you need to adjust this parameter: dynamic_axes = {'images': {0: 'batch_size', 1: 'channel', 2: "height" , 3: 'width'}, 'output': {0: 'batch'}}, images are the inputsname set before, 0: 'batch_size', 1: 'channel', 2: "height", 3: 'width 'means batch, channel, h, w are all dynamic dimensions, that is, as long as your model can process it, it can have a normal return value. For example: your image input cannot be input strictly according to 640, 640, and the output is the same
• Opset_version is to convert the version of onnx

ONNX model inference

import nmupy
import onnxruntime 
device = torch.device("cpu")
#onnx路径
model_path = "weights/yolov5s.onnx"
#数据前处理
img = numpy.transpose(image_tensor,(2,0,1))
img = numpy.expand_dims(img, 0) #添加一个维度 就是batch维度
img = img.astype(numpy.float32)#格式转成float32
img /= 255
#加载onnx模型
ort_session = onnxruntime.InferenceSession(model_path, providers=device)
#调用onnxruntime run函数进行模型推理
outputs = ort_session.run(
        None,
        {"images": img},
    )
re=numpy.array(outputs)

The reasoning of onnx is similar to that of weight file reasoning. You can choose cpu loading or cuda, and the process of data preprocessing is processed according to your own model. The key point is that the final format must be converted to float32, otherwise an error will be reported. Then the model loads onnxruntime.InferenceSession() and then calls the run function. It is worth noting here that "images" should be the inputname set when the model is converted, and the final output is a list type.