ONNXRuntime概览

深度学习

---

---

一、ONNXRuntime是什么？

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

虽然大家用ONNX时更多的是作为一个中间表示，从pytorch转到onnx后直接喂到TensorRT或MNN等各种后端框架了= =，但这并不能否认ONNXRuntime是一款非常优秀的推理框架（微软出品，必属精品）。而且由于其自身只包含推理功能（1.2版本，最新的ONNXRuntime甚至已经可以训练，可见微软在其上面的野心还是有的），对比主流框架源码看起来没有那么复杂难懂，通过阅读其源码可以非常清晰的理解深度学习框架的一些核心功能原理（op注册，内存管理，运行逻辑等）。接下来的一系列文章尝试对ONNXRuntime的源码进行阅读学习，对理解深度学习框架的工作原理还是很有帮助的。

二、使用介绍

1.引入库下面是使用ONNXRuntime的一个简单例子：

代码如下（示例）：https://github.com/skymuyu/coordconv-pytorch-mnist

# coding: utf-8
import onnxruntime as rt
import numpy
from PIL import Image


def softmax(x):
    """
    Compute softmax values for each sets of scores in x.

    Rows are scores for each class.
    Columns are predictions (samples).
    """
    scoreMatExp = numpy.exp(numpy.asarray(x))
    return scoreMatExp / scoreMatExp.sum(0)


def main():

    img = Image.open('../data/0.pgm')
    #print(img)
    x = numpy.array(img)
    #print(x[27][27])
    x = ((1-(x/255.0))- 0.1307) / 0.3081
    x = numpy.repeat(numpy.expand_dims(numpy.expand_dims(x, 0), 0), 6000, axis=0)
    #print(x[0])
    sess = rt.InferenceSession("../models/output.onnx")
    input_name = sess.get_inputs()[0].name
    print(input_name)
    output_name = sess.get_outputs()[0].name
    print(output_name)
    pred = sess.run([output_name], {input_name: x.astype(numpy.float32)})[0]
    print(pred[0])
    print(list(map(lambda x: "%.4f" % x, softmax(pred[0]))))

if __name__ == "__main__":
    main()

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

第一阶段 Session构造

构造阶段即创建一个InferenceSession对象。在python前端构建Session对象时，python端会通过http://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进行计算。

总结

以上就是ONNXRuntime的整体概览，由于ONNXRuntime仅仅是一个推理引擎，运行过程是线性的，逻辑也比较简单易懂。接下来的几篇文章会针对ONNXRuntime的一些核心功能（如内存管理，OpKernel注册，模型切分等）的代码进行详细剖析。