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