介绍

ONNX(Open Neural Network Exchange)是神经网络的一种交换文件格式，它既能够储存网络结构也能够储存网络的所有参数，它可以被认为是目前生产环境下部署神经网络的主流文件格式。ONNX文件描述了网络的所有计算算过程，这就意味着实际上它的作用类似传统意义上的二进制文件，是“可执行”的代码，其运行环境就是本文描述的部署方法。下面列举导出方法，原生C/C++ Java环境的部署和运行以及低级嵌入式平台部署。

导出与查看

在Pytorch下导出模型及权重到ONNX文件的代码示例如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

import torch
import onnxsim
# model为已经加载参数的模型变量，在定义模型结构后可以采用以下代码加载
# model.load_state_dict(torch.load('model_weights.pth'))
# model.eval()
# torch_input为模型输入参数大小，对于图片一般格式为N(Batch大小)*C(图像通道，灰度图像为1)*H(图像高)*W(图像宽)
# torch_input张量至少有两维，第一维必须是Batch大小
torch_input = torch.randn(1,1,28,28) 
torch.onnx.export(model,torch_input,"model.onnx", export_params=True
	,input_names=["input0"],output_names=["output0"],verbose=False)

# 重新加载并检查模型，自动化简模型
model_onnx = onnx.load("model.onnx")
onnx.checker.check_model(model_onnx)
model_onnx, check = onnxsim.simplify(model_onnx,
                dynamic_input_shape=False,
                input_shapes=None)
onnx.save(model_onnx, "model.onnx")

其中 torch.onnx.export 中的模型输入张量至少有两维，第一维必须是Batch大小，即为上述代码中的 torch_input 。这里的模型输入为图像，其一般格式为N(Batch大小)C(图像通道，灰度图像为1)H(图像高)W(图像宽)。最后化简模型一步是可选的。
模型的查看可以使用 netron 工具可视化查看网络结构

原生C/C++部署

• OnnxRuntime 微软官方ONNX的CPU调用库，支持Windows下C#/C/C++，当然也有Linux的支持但是需要 手动下载 。OnnxRuntime的 另外一个版本 是支持CUDA/TensorRT的GPU调用库
• OpenVINO 针对Intel平台优化的ONNX调用库
• OpenCV 虽然OpenCV是做图像处理的，但是它的dnn模块中也有加载ONNX模型调用的功能。

Java部署

• OnnxRuntime 微软官方ONNX的CPU调用库也是支持Java调用的，本质上是在Java内调用C/C++代码
• DJL(Deep Java Library) Java主流深度学习框架

低级嵌入式平台部署

• ESP32平台部署 部署过程比较复杂，需要先用 TVM 将模型参数量化，并且将一些算子转换为ESP32平台支持的加速算子，最后将模型生成为C代码
• STM32平台部署 用STM32CubeMX + X-CUBE-AI完成部署，部署平台至少Cortex-M4以上

参考

• Export a PyTorch model to ONNX
• Install ONNX Runtime (ORT)
• ONNX Runtime Execution Providers
• DJL - ONNX Runtime engine implementation