ONNX & ONNX Runtime 调研

0. 前言

• 最近要做一些模型推理优化的工作，所以要稍微多学习一点。
• ONNX其实一直都知道，统一的模型格式，使得模型在各个框架、平台间迁移更方便。

1. ONNX

• 官方资料：官网，官方Repo，还有各种repo
• 其他参考资料：ONNX–跨框架的模型中间表达框架

1.1. 基本情况

• ONNX是Open Neural Network Exchange的缩写，是与框架无关的模型表达方式。
• ONNX本身不提供太多功能，但各类工具都围绕着ONNX展开。
• ONNX相关工具（从官网的划分方式来看）
  • Build Models，即ONNX转换工具
    • 将各类深度学习框架训练得到的模型转换为统一的形式，比如PyTorch转ONNX、TensorFlow转ONNX等。
  • Deploy Model，即部署工具，后面单独介绍
  • Visualize，即可视化工具，如Neuron、Visual DL
• 已有的ONNX模型，可以到官方repo中找。
• 官方还提供了 tutorials，主要包括了
  • 模型转换（转换为ONNX或ONNX转换为别的）
  • 模型推理
  • 模型转换+部署
  • 模型量化
  • Serving（部署ONNX模型，这个倒是可以看一下）
  • 作为中介，将PyTorch模型转换为TF模型
  • 性能对比
  • 可视化模型

1.2. 部署工具概述

工具名称	CPU	GPU	Python	C++	备注
TensorRT	不支持	支持	支持	支持	GPU上性能最好
TVM	支持	支持	支持	支持	之前试用的时候感觉对ONNX的支持并不是特别好，需要比较多的自己改源码。不知道现在如何了
NCNN	支持	支持	不支持	支持	腾讯出品移动端推理优化工具，特点是对ARM CPU以及手机芯片优化做得比较好
ONNX Runtime	支持	支持	支持	支持	支持各种后端（比如TensorRT/OpenVINO等），感觉百搭，但不知道性能如何
OpenVINO	支持	支持	支持	支持	针对英特尔硬件的推理优化工具，如果是英特尔CPU，可能就用这个比较好

1.3. ONNX文档阅读

• 支持的OPS列表
• Python相关
  • examples与简单使用文档
  • 提供的功能包括模型导入、模型保存、检查模型、优化模型（这个优化应该是ops层面的）、提取子模型（原始模型的一部分）、polish model（不知道该怎么翻译，反正就是包括checker/optimizer/shape inference这几部）、版本转换（多个ONNX版本之间的转换）。

1.4. 其他小工具

• onnx-simplifier：模型结构优化工具
• onnx/optimizer：上一个工具的基础，即1.3中提到的模型优化，应该是做计算图优化
• onnx/backend-scoreboard：用来表示backend对onnx标准的compaliance（兼容性？）
  • 但实际没几个backend做了这个东西。

2. ONNX Runtime

• 官方资料：官网，官方repo，官方samples，官方文档

2.1. 基本情况

• 不仅支持推理，还支持训练（但不知道为啥要用这个训练……）
• 支持各种后端

2.2. 各种后端

• 需要注意的是，下图中这么多排列组合，并不是所有都支持。
  • 一部分直接支持
  • 一部分需要自己编译源码
  • 还有一部分没有测试过
  • 
• Python API 介绍
  • pip install onnxruntime
    • Windows + Default CPU 时，在X64中可直接使用。
    • Linux + Default CPU 时，在X64/ARM64中可直接使用。
  • pip install onnxruntime-gpu
    • Windows/Linux + CUDA时，可在X64下使用。
• TensorRT相关
  • 都需要从源码编译
  • 要求的TensorRT版本和CUDA版本都比较高，不知道低版本能不能用。
  • 不知道性能跟直接用TensorRT比如何。

2.3. Python API 的基本使用

• 官方文档
• 使用有点类似于TF

import numpy
import onnxruntime as rt

sess = rt.InferenceSession("logreg_iris.onnx")
input_name = sess.get_inputs()[0].name
pred_onx = sess.run(None, {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)