ONNX作为模型部署的中间变量，无论是tensorRT还是tfsavemodel都需要先将pytorch或者tensorflow模型转换为ONNX，ONNX作为中间变量也可以通过onnxruntime进行推理。

ONNX模型转换

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应该为事先训练好的模型.pth结尾或者pt结尾的文件
model的获取需要更改，也就是你的模型类的实例化这个自行更改
由于训练时权重保存的方式是有所不同的，分为两种，一种是只保留权重，那么就需要利用load_state_dict将读取到权重文件加载给模型。另一种则是将权重和模型一起进行保存，这样模型获取这个过程也进行了省略，可以直接model=torch.load(weight)这样就直接获得了加载好权重的模型。
model.eval()做推理的时候一定要调用eval。
dummy_input做转换onnx时候要先定一个输入shape的张量，如果下面dynamic_axes这个参数设置为None，那吗这个张量一定要和没次推理模型接受的shape完全一样。
output_name需要制定onnx文件保存的路径和名字你可以自行定义例如/user/onnx/model.onnx
torch.onnx.export()进行torch到onnx的转换除了上面提到的一些参数外，还需要定一个input_names与output_names
模型转换时可以设置动态输入，也就是你的输入可以没次都不一样，需要调整这个参数：dynamic_axes={'images': {0: 'batch_size', 1: 'channel', 2: "height", 3: 'width'},'output': {0: 'batch'}}，images就是之前设置的inputsname，0: 'batch_size', 1: 'channel', 2: "height", 3: 'width'就代表batch，channel，h，w这几个维度都是动态的就是只要你的模型能进行处理，就可以有正常的返回值。举个例子：就是你图像输入不能严格按照640，640进行输入，output同理
opset_version就是转换onnx的版本

ONNX模型推理

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)

onnx的推理和权重文件推理的方式类似，可以选择cpu加载或者cuda，数据前处理的过程就根据你自己的模型进行处理，重点说一下最后格式要转换成float32，要不然会报错。然后模型加载onnxruntime.InferenceSession()后然后调用run函数，这里值得注意的是“images”要为模型转换时设置的inputname，最后的输出为list类型。

ONNX模型转换与onnxruntime推理

ONNX模型转换

ONNX模型推理

猜你喜欢