pytorch: 转onnx模型

摘要： onnx(Open Neural Network Exchange)主要用于部署，训练一般使用pytorch和tensorflow，等训练完成，我们需要模型转成onnx，用于部署在云或者边缘计算上。而为什么要要转成onnx模型呢，主要是因为onnx没有训练，只有推理，速度很快，而且目前大多数芯片都适配onnx模型，相当于一个通用莫模型，易部署，而且速度快。

1. 导出onnx模型

使用torch.onnx.export()可以把pytoch模型转成onnx模型，其中onnx.export()函数是自带的函数。

import torch 

batch_size = 1  
input = torch.rand(batch_size, 3, 256, 256)  # 初始化输入
input = input.cuda()
model = model.cuda()  # 加载模型
model.eval()  # 把模型设置成推理模式
output = model(input)  # 推理
print("Beginning to export onnx ...")
torch.onnx.export(model,  # 要转换的模型
                  input,  # 输入
                  "onnx/model.onnx", # 到处的模型，以onnx格式保存
                  export_params=True,  # 训练参数是否和模型一同导出，一般设置为true
                  opset_version=12,  # 导出onnx模型的版本
                  verbose=False,  # 是否打印参数
                  input_names=["input"],  # 输如名称
                  output_names=["output"],  # 输出名称
                  dynamic_axes={
                                "input": {0: "batch_size"},
                                "output": {0: "batch_size"}
                                }  # 把输入的批量维度设置为动态
                  )
print("Exporting onnx successfully!")

最后，在文件夹onnx下会生成一个model.onnx文件。

2. 验证onnx模型

安装onnx模块：

pip install onnx

验证onnx模型是否成功导出：

import onnx

model_onnx = onnx.load("onnx/model.onnx")  # 加载onnx模型
onnx.checker.check_model(model_onnx)  # 验证onnx模型是否成功导出

如果没有报错，表示导出成功

3. 运行onnx模型

安装onnxruntime，如果支持cuda，则安装onnxruntime-gpu

pip install onnxruntime-gpu

运行onnx模型

import onnxruntime
import numpy as np

# 创建会话
session = onnxruntime.InferenceSession("onnx/model.onnx",
                        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
x = np.random.randn(1, 3, 256, 256)
ort_input = {session.get_inputs()[0].name: x.astype(np.float32)}
ort_output = session.run(None, ort_input)