ONNX Runtime 源码阅读:引擎运行过程总览

推荐阅读
ONNX Runtime 源码阅读:模型推理过程概览
ONNX Runtime 源码阅读:模型结点串行执行顺序的确定
ONNX Runtime 源码阅读:模型分区的实现

之前的文章中,我们一直从源码细节的角度去看某个功能的实现。总会有一种一叶障目不见泰山的感觉,对ONNX Runtime缺乏一个比较总体的认识。通过前面的几篇文章,在看细节源码的过程中,脑海渐渐浮现出了ONNX Runtime总体的轮廓,但是并不是很清晰。今天打算暂时跳出细节,站在高处,看看ONNX Runtime它的筋骨、它的脉络,也更能方便以后进行局部分析。
下图展现的是ONNX Runtime主要的一些类以及它们之间的关系,这里只展现了比较重要的一些类而不是全部。通过UML图可以看到,InferenceSession, SessionState这两个类占据着举足轻重的地位。如果要打个比方,那么InferenceSession就是一个运筹帷幄之中的统帅,负责统筹全局、制定计划;而SessionState就相当于保障物资的后勤保障系统。所谓兵马未动粮草先行,在小兵们(OpKernel)冲锋陷阵(推理)之前,统帅一定要准备好各种粮草辎重、武器装备并制定战术,不仅让士兵吃饱、装备精良,还要尽量发挥各个兵种的优势,该远程消耗就让弓箭手上、该攻城就让步兵来,不然让骑兵去打水战注定是悲剧。
下面就让我们站在上帝的视角,看一看一场战斗是怎么打赢的。
overview.jpg

之前多次说过,模型推理的过程可以简单分为三步:实例化、初始化、推理。今天就这上面这张图,分析一下整个推理过程。

实例化

实例化的时候,InferenceSession会构建起一个框架。当实例化结束,InferenceSession已经持有SessionState, KernelRegistryManager, ExecutionProviders等的引用,只不过他们也都只是个空架子,里面并没有什么实质性内容。紧接着,就进入了初始化阶段。
初始化阶段是模型推理的重头戏,它最终让一个静静躺在磁盘中的模型文件中的描述内容再内存中一一建立了起来。

初始化

模型加载

首先,InferenceSession将模型文件从磁盘读入并解析。如果模型内容已经存在内存之中,就直接使用这块内存;如果模型还不在内存,就通过一个Load()方法去加载。值得注意的是,虽然解析到的模型可以通过一个Protobuffer编译器编译出来的model->MainGraph可以直接访问到模型的内容,但是ONNXRuntime还是在这之上添加了一层抽象层,这个抽象层就是按照ONNX标准来的。主要目的就是使得整个引擎和模型文件的格式隔离开,引擎最终与模型文件个格式解耦,也就是和模型的解析过程解耦。虽然现在ONNX用的是Protobuffer,但是有了则层抽象层,后面如果换成了FlatBuffer或者其他的格式,也不会对整个框架有影响。并且,最终Load()方法会委托用户自定义的一个loader去执行具体将模型加载到内存的操作,这就让从哪里加载模型获得了极大的自由。突然想到Java的类加载器,只要你最终加载到的模型符合标准,并不管你是从本地磁盘还是网络获取到的。
模型加载并解析之后,得到的是一个原始的模型在内存中的表示。原始的意思就是,整个数据结构是直接通过Protobuffer编译器编译得到的,后续还需要对她进行一层封装处理。都说巧妇难为无米之炊,现在米已经有了,就看怎么做好这顿饭。

注册Provider以及Kernel

在实例化阶段,已经生成了KernelRegistryManager的实例。此时,InferenceSession将会实例化所有可用Provider,并将他们保存在ExecutionProviders当中。之后,KernelRegistryManager会通过ExecutionProviders拿到InferenceSession传递过来的Provider,并委托他们提供自己所支持的所有OpKernel的信息。这些信息并不是OpKernel的实例,而是KernelRegisty的实例。KernelRegisty总持有生成每一OpKernel的实例所需要的信息和方法,这些信息和方法通过一个结构体KernelCreateInfo保存着。当这一步完成,有关锅碗瓢盆也准备停当。

模型分区

InferenceSession委托GraphPartitioner去做模型分区工作。GraphPartitioner初始化的时候InferenceSession会给到它ExecutionProvidersKernelRegistryManager。具体分区过程在ONNX Runtime 源码阅读:模型分区的实现中有介绍。

模型节点执行顺序的确定

模型加载并解析之后,所的到的知识原始的模型结构,ONNX Runtime需要对它按照ONNX的标准进行转换和封装,包括将所有原始节点编号、确定执行顺序等,然后存入Graph当中,最后Graph被存放于SessionState当中。具体执行顺序的确定的过程在ONNX Runtime 源码阅读:模型结点串行执行顺序的确定中有介绍。

实例化Kernel

Kernel的实例化过程中,InferenceSessionKernelRegistryManager传递给了SessionState.CreateKernels()。反过来,SessionState.CreateKernels()拿到KernelRegistryManager后委托它通过来实例化。最终通过层层委托,最终是KernelCreateInfoKernelCreateFn根据KernelDef最终生成了OpKernel的实例。这部分就留着以后细看吧。

推理

推理,也就是依次调用OpKernel.Compute(),利用面向对象的多态机制,此时基类OpKernel的引用绑定的是每一个具体的子类,因此会调用到子类重写的Compute()方法。

另外还有一大块内容:运行过程中的内存分配,我们从来没有触碰过。可能对于运行在服务器端的引擎还没那么敏感,但是内存对于移动设备那可真是寸土寸金。这些都值得一看,留坑后续慢慢填了。

公众号二维码

首发于个人微信公众号TensorBoy。微信扫描上方二维码或者微信搜索TensorBoy并关注,及时获取更多最新文章!
C++ | Python | 推理引擎 | AI框架源码,有一起玩耍的么?

Resources

https://github.com/Microsoft/onnxruntime
https://github.com/zmychou/onnx-runtime-code-reading

发布了45 篇原创文章 · 获赞 4 · 访问量 1万+

猜你喜欢

转载自blog.csdn.net/ZM_Yang/article/details/105158051