简单的将pytorch模型部署到onnx

1. 创建一个pytorch模型

这里我用的U2Net，直接加载好训练出的权重

model = U2Net(class_nums=4)
model.load_state_dict(torch.load(checkpoint_path))

2. 将pytorch模型转成onnx格式

x = torcg.randn(1,3,512,512)

with torch.no_grad():
    torch.onnx.export(
            model,
            x,
            "onnx_file_name.onnx",
            opset_version=11,
            input_names=['input'],
            output_names=['output']
)   

torch.onnx.export()函数用于将pytorch模型转成onnx格式的函数，其中的参数如下

• model：需要转换的模型，加载好权重
• args：模型的任意一组输入，注意维度
• f：导出的onnx模型的文件名
• opset_version：表示ONNX算子集的版本，随着深度学习的发展，新算子会不断诞生，算子集就需要扩充，截至20230824，官网已经发布20个版本
• input_names：输入tensor的名称
• output_names：输出tensor的名称

转换成功会在指定的路径f下生成onnx模型文件

可用onnx中的函数验证模型文件是否正确

import onnx

onnx_model = onnx.load("onnx_file_name.onnx")
try:
    onnx.checker.check_model(onnx_model)
except Exception as ex:
    print(f'ERROR: {ex}')
else:
    print(f'model correct')

 可以使用netron查看转换的onnx模型的详细内容  链接 https://netron.app/

 点击某个算子节点，可以查看算子的具体信息，每个算子记录了算子的属性、图结构和权重三类信息

• 算子属性信息(attributes)，对于卷积来说，算子属性包含了卷积核的大小、步长等属性
• 图结构信息(node properties)，算子节点在计算图中的名称、邻边的信息，如上图节点名为conv2，输入数据叫683，权重叫1324
• 权重信息指算子存储的网络权重信息，点击w后的+号可以查看权重信息的具体内容

3. 推理引擎ONNXRuntime

使用onnxruntime可以运行输出的onnx模型，代码如下

input_img 是一个与onnx模型输入维度一致的归一化的np.array

model = onnxruntime.InferenceSession(onnx_file)
inputs = {'input':input_img}
output = model.run(['output'],inputs)[0]

onnxruntime.InferenceSession用于获取一个ONNX Runtime的推理器，其参数为用于推理的ONNX模型文件

run方法用于模型推理，第一个参数为输出张量名的列表，第二个参数为输入值的字典，key为导出模型时设置的张量名，value为输入张量