【AI芯片】ONNXRuntime

ONNXRuntime是微软推出的一款推理框架，用户可以非常便利的用其运行一个onnx模型。ONNXRuntime支持多种运行后端包括CPU，GPU，TensorRT，DML等。可以说ONNXRuntime是对ONNX模型最原生的支持。

虽然大家用ONNX时更多的是作为一个中间表示，从pytorch转到onnx后直接喂到TensorRT或MNN等各种后端框架，但这并不能否认ONNXRuntime是一款非常优秀的推理框架。而且由于其自身只包含推理功能（最新的ONNXRuntime甚至已经可以训练），通过阅读其源码可以解深度学习框架的一些核心功能原理（op注册，内存管理，运行逻辑等）

import numpy as np
import onnx
import onnxruntime as ort

image = cv2.imread("image.jpg")
image = np.expand_dims(image, axis=0)

onnx_model = onnx.load_model("resnet18.onnx")
sess = ort.InferenceSession(onnx_model.SerializeToString())
sess.set_providers(['CPUExecutionProvider'])
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

output = sess.run([output_name], {input_name : image_data})
prob = np.squeeze(output[0])
print("predicting label:", np.argmax(prob))

总体来看，整个ONNXRuntime的运行可以分为三个阶段，Session构造，模型加载与初始化和运行。和其他所有主流框架相同，ONNXRuntime最常用的语言是python，而实际负责执行框架运行的则是C++。

第一阶段 Session构造

构造阶段即创建一个InferenceSession对象。在python前端构建Session对象时，python端会通过onnxruntime_pybind_state.cc调用C++中的InferenceSession类构造函数，得到一个InferenceSession对象。

InferenceSession构造阶段会进行各个成员的初始化，成员包括负责OpKernel管理的KernelRegistryManager对象，持有Session配置信息的SessionOptions对象，负责图分割的GraphTransformerManager，负责log管理的LoggingManager等。当然，这个时候InferenceSession就是一个空壳子，只完成了对成员对象的初始构建。

第二阶段 模型加载与初始化

在完成InferenceSession对象的构造后，会将onnx模型加载到InferenceSession中并进行进一步的初始化。

1 . 模型加载
模型加载时，会在C++后端会调用对应的Load()函数，InferenceSession一共提供了8种Load函数。包读从url，ModelProto，void* model data，model istream等读取ModelProto。InferenceSession会对ModelProto进行解析然后持有其对应的Model成员。
2 . Providers注册
在Load函数结束后，InferenceSession会调用两个函数：RegisterExecutionProviders()和sess->Initialize();

RegisterExecutionProviders函数会完成ExecutionProvider的注册工作。这里解释一下ExecutionProvider，ONNXRuntime用Provider表示不同的运行设备比如CUDAProvider等。目前ONNXRuntimev1.0支持了包括CPU，CUDA，TensorRT，MKL等七种Providers。通过调用sess->RegisterExecutionProvider()函数，InferenceSession通过一个list持有当前运行环境中支持的ExecutionProviders。

3 . InferenceSession初始化
即sess->Initialize()，这时InferenceSession会根据自身持有的model和execution providers进行进一步的初始化（在第一阶段Session构造时仅仅持有了空壳子成员变量）。该步骤是InferenceSession初始化的核心，一系列核心操作如内存分配，model partition，kernel注册等都会在这个阶段完成。

首先，session会根据level注册 graph optimization transformers，并通过GraphTransformerManager成员进行持有。
接下来session会进行OpKernel注册，OpKernel即定义的各个node对应在不同运行设备上的计算逻辑。这个过程会将持有的各个ExecutionProvider上定义的所有node对应的Kernel注册到session中，session通过KernelRegistryManager成员进行持有和管理。
然后session会对Graph进行图变换，包括插入copy节点，cast节点等。
接下来是model partition，也就是根运行设备对graph进行切分，决定每个node运行在哪个provider上。
最后，为每个node创建ExecutePlan，运行计划主要包含了各个op的执行顺序，内存申请管理，内存复用管理等操作。

第三阶段 模型运行

模型运行即InferenceSession每次读入一个batch的数据并进行计算得到模型的最终输出。然而其实绝大多数的工作早已经在InferenceSession初始化阶段完成。run阶段主要是顺序调用各个node的对应OpKernel进行计算。