Page 35 - 《软件学报》2020年第10期

徐梦炜 等:面向移动终端智能的自治学习系统 3011

些框架实现了各种算子的前向/后向运算,以第三方库的方式加载进内存,然后读入模型与数据,进行运算;另一
种方式是近些年兴起的以 TVM [20] 为代表的神经网络编译器,这些编译器为特定的模型和硬件平台生成一个可
执行文件,可直接在该平台上部署运行.例如,TVM 沿用并拓展了 Halide [21] 的 compute-schedule 这一概念,通过定
义 schedule 来确定代码的优化过程,通过做自动代码生成来实现在各种设备上的部署深度学习模型.相较于通
用的机器学习框架,神经网络编译器具有以下几个优点:(1) 可以根据硬件和算子的描述,自动生成底层的运算
代码,具有更强的扩展性;(2) 可以自动搜索在特定硬件上的实现方式(包括 memory layout,SIMD),以达到最优
的运算效率.在自治学习场景下,引入模型训练编译器可以有效地减少终端设备上的计算开销.
然而,已有的神经网络编译器只服务于模型推断的场景,无法用于模型训练.为此,我们扩展了原有的 TVM
框架,以实现移动终端上的自治学习.
如图 4 所示,为了编译生成一个可用于训练的可执行文件,我们将用户输入的网络结构和损失函数通过编
译得到 3 个网络.
(1) 可运行的前向网络,该网络用于前向计算;
(2) 可运行的 loss 网络,该网络与可运行的前向网络的区别在于增加了损失函数,通过这一个网络可以得
到相关的训练参数,从而确定需要更新的参数;
(3) 可运行的更新网络,该网络用于更新网络的相关参数,根据输入参数和计算出来的相关梯度,利用 SGD
等算法将参数进行更新,从而完成对计算图的参数更新.
生成以上 3 个不同功能网络的目的是为了更好地利用 TVM 模块化特性,减少对 TVM 代码库的修改.
最后,我们补充了 TVM 中缺失的各个算子的后向传播,从而对于任意模型结构,可以编译得到在移动端执
行的二进制代码,最后在移动设备上运行更新网络即可以实现训练过程.由于模型训练中的反向传播与模型推
断在底层实现上基本类似,例如最核心的操作都是围绕矩阵乘法展开的,因此在编译出可执行代码的过程中,我
们沿用了推断过程中已有 schedule 的设计思路,包括搜索空间与搜索算法.更具体地,我们针对上述第 3 个可运
行的更新网络,增加了针对反向传播算子(如 conv2d_grad、dense_grad 等)的 schedule 的搜索空间,以 TVM 现有
的 template 形式实现;对于搜索算法,我们沿用了 TVM 中基于机器学习的 cost model 方法来预测 schedule 的运
行时间,并在实际编译的过程中不断更新这个 cost model,最终使搜索过程快速达到收敛.

Fig.4 A deep learning compiler for autonomous learning
图 4 面向自治学习的模型编译框架

此外,针对不同的终端平台,TVM 内部实现了很多与硬件相关的 schedule 原语,我们利用这些原语来设计不
同终端平台的代码搜索空间,经过自动调优的不断测试,最终得到针对后端硬件优化后的代码.本文扩展后的
TVM 框架在生成训练代码时,同样需要使用不同终端设备(或不同硬件 SoC)进行各自的调优,调优过程中需要
真实设备提供运行时服务,以获取搜索过程中遍历节点的运行效率,作为搜索过程的反馈.对于每一台设备或硬
件(本文实验中使用了两种设备),都需要进行以上搜索调优过程,时间成本开销较大.在未来工作中,我们将尝试
使用更加有效的 cost model [22] 来加速编译过程.