深度学习模型部署框架之ONNX Runtime

0 前言

本文是深度学习模型部署框架的第一篇，更多部署框架可以查看本专栏的其他文章。

1 简介

概念：onnx一种开源模型格式，专注于推理所需的能力。同时，它还是一个中间深度学习框架，用于连接不同深度学习框架之间的转换。
优势：跨平台兼容性好；性能优化；支持多种框架。
坑点：onnx只是一个格式，就和json一样，只要满足规则，就是合法的。因此单纯从pytorch转成onnx格式很简单，但是不同后端框架接受的onnx是不一样的。比如pytorch自带的torch.onnx.export转换得到的onnx，onnxruntime 需要的onnx，TensorRT需要的onnx都是不同的，因此这才是坑的来源。

2 模型准备

2.1 导出模型为onnx格式

• 使用模型转换工具，torch.onnx.export和

# pt->onnx
# 示例数据
dummy_input = torch.randn(1,1, 16000*5)  # 只有第三个维度可能会变化，前两个都不会变
# 导出模型为 ONNX 格式
torch.onnx.export(
    model,                   # PyTorch 模型
    dummy_input,             # 示例输入
    "denoiser_model_12_3_3_dongtai_5s.onnx",  # ONNX 文件的保存路径
    input_names=['input'],   # 输入层的名字
    output_names=['output'], # 输出层的名字
    dynamic_axes={
        'input': {2: 'audio_length'},  # 设置动态维度,如果输入是三维的，那就是第三个维度是动态的，否则是第二个维度是动态的
        'output': {2: 'audio_length'}
    },  # 可选，设置哪些维度是动态的
    opset_version=12        # ONNX opset 版本，建议使用 11 或更高
)

# tf->onnx
python -m tf2onnx.convert
    --saved-model SOURCE_SAVED_MODEL_PATH |
    --checkpoint SOURCE_CHECKPOINT_METAFILE_PATH |
    --tflite TFLITE_MODEL_PATH |
    --tfjs TFJS_MODEL_PATH | 
    --input | --graphdef SOURCE_GRAPHDEF_PB
    --output TARGET_ONNX_MODEL

2.2 导出时注意事项

• 动态形状的支持，对于音频来说基本上都需要动态输入，这需要在模型转换的时候自定义
• 模型的opset版本

2.3 onnx模型检查

• onnx.checker验证模型
• Netron可视化模型
• 检验onnx的输出与原始模型的输出是否一致，或输出的误差是否可以容忍

3 模型优化

3.1 量化

• 量化分为动态量化和静态量化，常用的方法是动态量化，相对来说简单一些。
• 动态量化：对于动态量化，缩放因子和零点是在推理时计算的，并且特定用于每次激活，因此它们更准确，但引入了额外的计算开销。

from onnxruntime.quantization import QuantType, quantize_dynamic
 
# 模型路径
model_fp32 = 'models/denoiser.onnx'
model_quant_dynamic = 'models/denoiser_quant_dynamic.onnx'
 
# 动态量化
quantize_dynamic(
    model_input=model_fp32, # 输入模型
    model_output=model_quant_dynamic, # 输出模型
    weight_type=QuantType.QUInt8, # 参数类型 Int8 / UInt8
    optimize_model=True # 是否优化模型
)

• 静态量化：对于静态量化，使用校准数据集离线计算，所有激活都具有相同的缩放因子和零点。这点类似于tfl的int8量化。
• 通常来说，对 RNN 和基于 Transformer 的模型使用动态量化，对 CNN 模型使用静态量化

3.2 算子融合

其实就是使用onnx-simplifier简化模型，例如 Conv+BatchNorm 融合。

pip install onnx-simplifier
python -m onnxsim model.onnx model_optimized.onnx

3.3 硬件加速

ONNX Runtime 支持多种硬件加速器，如 CUDA 加速。

import onnxruntime as ort

# 加载模型
sess_options = ort.SessionOptions()
sess = ort.InferenceSession("model_quantized.onnx", sess_options, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

3.4 多线程

ONNX Runtime 支持多线程执行，可以通过设置环境变量或通过 API 设置线程数来控制。

sess_options.intra_op_num_threads = 4  # 设置线程数量
sess = ort.InferenceSession("model_quantized.onnx", sess_options)

4 利用onnx模型推理的demo

import onnxruntime as rt
import onnx
import numpy as np
import time
 
ONNX_Model_Path = 'path_to_onnx_model'
 
img = np.ones((1, 3, 0, 0)).astype(np.float32)
 
# 加载模型，传入模型路径
model = onnx.load_model(ONNX_Model_Path)
 
# 创建一个SessionOptions对象
rtconfig = rt.SessionOptions()
 
# 设置CPU线程数为4
cpu_num_thread = 4
 
# 设置执行模式为ORT_SEQUENTIAL(即顺序执行)
rtconfig.intra_op_num_threads = cpu_num_thread
rtconfig.execution_mode = rt.ExecutionMode.ORT_SEQUENTIAL
 
# 设置使用的ExecutionProvider为CPUExecutionProvider
providers = ['CPUExecutionProvider']
 
# 创建一个InferenceSession对象
sess = rt.InferenceSession(model.SerializeToString(), providers=providers, sess_options=rtconfig)
 
# 模型的输入和输出节点名，可以通过netron查看
input_name = 'audio'
outputs_name = ['output']
 
# 模型推理:模型输出节点名，模型输入节点名，输入数据
 
net_outs = sess.run(outputs_name, {input_name: img})
result = np.array(net_outs)

5 参考文档

• 文档：https://onnx.ai/onnx
• onnx model：https://github.com/onnx/models
• onnx tutorials：https://github.com/onnx/tutorials
• 参考示例的demo：https://onnxruntime.ai/docs/build/android.html、 https://github.com/microsoft/onnxruntime-inference-examples

总结

本文介绍了onnxruntime的在应用方面的一些经验总结，更多内容参阅文档。