对于在移动端和IoT设备上部署深度学习模型,通用、高性能的移动端推理框架起到十分关键的作用。在使用mnn过程中,自己越发对其内部实现框架感兴趣,因此对MNN中推理部分代码设计思路做进一步阅读和整理。
1.简介
MNN 是阿里巴巴开源的一个轻量级的深度学习端侧推理引擎,旨在解决深度神经网络模型在端侧推理运行问题,包含深度神经网络模型的优化、转换和推理,具有高通用性和高性能的特点,支持多种训练框架的模型,常用深度学习算子,多种系统,以及卷积计算方式计算优化等。
上图是官方的代码模块图,从中可以看出MNN包含Converter和Interpreter两部分,其中Converter负责模型转换、模型优化; Interpreter负责推理。Interpreter包含Engine和Backends组成,其中Engine负责模型加载和调度,Backends负责内存分配和计算。接下来着重分析下Interpreter部分。
2. 推理部分
首先从整个推理使用流程出发,查看整体设计思想,进一步详细看不同Backends和深度学习算子的实现模式。
2.1 使用流程
1)调用createFromFile, 创建Interpreter对象;
2) 调用Interpreter对象的成员函数createSession, 创建会话session;
3) 调用Interpreter对象的成员函数getSessionInput获取指定会话下的输入Tensor;
4)调用resizeTensor, resizeSession 进行初始的内存预分配;
5)调用ImageProcess引擎,加载输入数据到输入Tensor,同时可以进行相关图像预处理操作,比如亮度上减均值除方差、图像缩放,rgba2bgr图像channel转换,nchw格式转换等;
6) 运行runSession,执行推理;
7)运行getSessionOutput获取指定session输出;
8)跨设备(CPU和GPU)操作时,调用copyToHostTensor进行数据复制操作,并执行后处理操作;
9)调用Interpreter对象releaseModel, releaseSession释放资源。
2.2)整体框架
沿着推理流程出发,可以看到整体流程如下:
Interpreter采用单例模式实现,保证全局只有一个该类的实例;一个Interpreter对象可以拥有多个Session。Session记录执行推理的Backend信息(即CPU、GPU等),和多条计算路径;Pipeline类记录了每条计算路径,包含执行算子的记录向量;Unit类表示单一计算op,有一个指向Execution的指针的成员变量。Execution类为抽象接口类,提供onCreate方法和OnExecute两个方法,其中onCreate创建指定Backend下特定计算op的Execution的派生类;onExecute执行计算。
当调用createSession()创建Session时,首先根据配置文件中硬件类型,创建指定Backend;其次创建表示执行路径Pipeline, 每个Pipeline遍历记录中的op创建Unit类,Unit类进一步创建Execution的派生类,负责实际的计算。
当调用runSession()执行推理操作时,依次调用Session类中成员函数run, 该函数遍历Pipeline,执行每个Pipeline中execute()成员函数;Pipeline类中execute函数,遍历执行Unit类中的excute函数;Unit类中excute函数委托Eexcution类中的onExcute函数负责指定硬件指定操作的计算。
2.3)支持不同Backend的设计思路
MNN目前支持cpu(CPU、arm82优化), GPU(opencl, GL, Vulkan, metal)几种类型,是基于经典的工厂模式实现的,如下图所示。每种Backend对应一个Creator,用来创建具体Backend。静态变量gExtraCreator记录指定设备对应的Creator,函数MNNInsertExtraBackendCreator负责注册新设备的Creator;函数MNNGetExtraBackendCreator负责获取指定设备的Creator;首次调用时完成支持Creator的一次注册,并且只注册一次。在此基础上,提供了BackendFactory类,供外部client调研。该类包含静态成员函数create, 获取指定设备的Creator,并创建Backend。
2.4)指定设备下op实现
在创建不同设备后,需要进一步创建op实现;不同设备对应一套算子,每个设备下算子仍由工厂模式实现。接下来以CPUBackend为例,查看具体实现。
在CPUBackend类中定义基类Creator,对应派生类,负责创建Execution的派生类,比如CPUArgmaxCreator 负责创建CPUArgMax。CPUBackend中静态成员函数addCreator负责注册op的Creator类;CPUBackend中onCreate函数,获取creator, 执行creator中的onCreate创建特定op的执行器。
3.总结
本文主要介绍MNN推理框架大体架构设计思路,首先从使用流程出发介绍整体架构,其次介绍了支持不同的Backend和不同op算子的工厂模式实现。