0%

ONNX模型使用与多环境部署

发表于 2024-06-19 分类于 Machine Learning

介绍

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

导出与查看

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

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以上

参考