2.2 挑战
许多现有的编译器(Chen et al., 2018; Vasilache et al., 2018)和框架(Paszke et al., 2019; Abadi et al., 2016; Jia et al., 2014)依赖基本的拓扑排序算法来调度图。虽然当前的方法可能足以在服务器级机器上运行传统网络,但这种方案可能不适合在资源受限的边缘设备上运行不规则连接的神经网络。这是因为,与运行规律拓扑的网络不同,运行不规则网络的内存占用范围会因为调度的不同而变化。例如,在一个代表性的边缘设备(SparkFun Edge: 250KB 的权重/激活内存和 60M MACs)上,图 3(b) 显示出仅有 4.1% 的调度勉强满足硬内存限制,而仅有 0.04% 的调度能够实现最优的峰值内存。实际上,这种限制将阻碍网络设计的多样性和创新探索。为了让边缘计算充分利用不规则连接的神经网络,应该缓解甚至消除这种限制。
2.3 设计目标
调度算法。为了解决这个问题,我们的工作旨在从搜索空间S 中找到一个节点调度 s* ,以最小化峰值内存占用 µpeak . S列举了图 G 中所有节点 v∈V 的所有可能排序,其中 V 是图中的所有节点集合。
调度的最直接方法是暴力法,该方法简单地枚举 S,并选择一个具有最小峰值内存占用的调度。尽管这种极端方法可能找到最优解,但由于其巨大的复杂性:Θ(|V |!) ,在时间上是过于昂贵的,其中 |V| 表示图中的节点数量。一种改进的方法是将搜索空间缩小,仅关注拓扑排序 ST ⊂ S。然而,这仍然会遭受复杂性的上限为 O(|V|!) 的问题(调度仅有 30 个节点的有向无环图(DAG)可能需要几天时间)。事实上,之前的研究(Bruno & Sethi, 1976; Bernstein et al., 1989)已经证明,对于 DAG 的最优调度是 NP 完全的。在另一种极端情况下,有启发式拓扑排序算法,如卡恩算法(Kahn, 1962),其复杂性为 O(|V|+|E|),其中 V 和 E 分别是节点和边的数量。然而,正如图 3 所示,该方法可能产生无法在目标硬件上运行的次优节点调度。为此,我们探索将动态编程与自适应软预算相结合的调度方法,以在保持图常量 s* ,的同时,达到最优解,而不会在时间上增加过多的开销。我们将在第 3.1 和 3.2 节中深入解释我们的算法。
图重写。任何调度算法,包括我们的方法,固有上都受到图拓扑的限制。因此,我们探索通过图重写(Plump, 1999)来转换搜索空间。图重写通常涉及用不同的模式替换图中的某些模式,以实现特定目标。对于计算数据流图,利用图中计算的分配性质、结合性质和交换性质,图重写可以在不改变语义的情况下,对某些目标带来显著的改进。例如,在一般程序中, Pi log x_i 可以表示为 P_{\text{odd}} \log x_i + P_{\text{even}} log x_i 或 log Q_i x_i ,而 ( x + x ) 可以转换为 ( x^2 ) 或 ( x << 1 )。同样,我们将这一见解引入神经网络,寻找一组可能的变换 X ,以将原始图 G 重写为一个新的图 G’ ,从而使我们的搜索空间 S 转变为具有更低峰值内存占用的空间:
我们确定了一组候选变换模式 x : g -> g’(其中 g ∈ G 且 g’ ∈ G’ ,构成 X )。在转换图时,我们的方法保持图的数学完整性,因此不是近似方法。我们将这种系统化的方法嵌入到提高峰值内存占用和搜索空间的过程中,称为身份图重写,并将在第 3.3 节中讨论该技术。
图 4. SERENITY 的整体工作流程:不规则连接神经网络的内存感知调度
3 SERENITY:不规则连接神经网络的内存感知调度
如第 2 节所述,目标是在执行不规则连接的神经网络时减少峰值内存占用。我们提出了 SERENITY,一种针对资源受限设备(例如边缘设备)的内存感知调度。图 4 概述了整体调度过程,突出了我们方法的主要贡献。SERENITY 的输入是一个不规则连接神经网络的图 G,实际上在调度过程中充当中间表示(IR)。我们为这个 IR 增加节点的元数据,例如操作类型、输入/输出边、输入/输出形状和内存成本。然后,图重写器将图 G 转换为 G’,以放宽内存密集模式的内存成本,目标是减少 G 的峰值内存占用µpeak。SERENITY 使用基于动态编程的调度器将图调度到最优调度 s^* 。然而,由于复杂性,调度可能较慢,因此我们通过利用分治法将搜索空间缩小,将图划分为多个子图。然后,我们为调度器增添自适应软预算,通过迅速的元搜索自适应地找到阈值预算,修剪次优路径,以加快调度过程。本节重点介绍 SERENITY 的创新点:基于动态编程的调度、分治法、自适应软预算和图重写,分别在第 3.1、3.2 和 3.3 节中详细解释。
3.1 基于动态编程的调度:实现最优峰值内存占用
我们对调度算法的目标是最小化峰值内存占用 ***µpeak (s,G) ***。如第 2.3 节所述,覆盖整个搜索空间 S 或所有拓扑排序子空间 ST ⊂ S 的递归算法需要极长的时间。这主要是由于对子问题的重复重新计算,使得算法的上限为 O(|V|!)。因此,我们利用动态编程 (Bellman, 1961; 1966; Held & Karp, 1962),其中包括一种记忆化方案,已被证明在通过重用子问题的解决方案来减少时间密集型算法的复杂性方面是有效的,同时仍然通过遍历整个搜索空间找到最优解。
识别签名以启用动态编程。将动态编程应用于新问题的第一步是表征最优解的结构:
( s^* = s^_n )(( s^n ) 是 n 个节点的最优解)。然后,需要识别子问题 ( s^_i ) 和原始问题 ( s^{i+1} ) 之间的递归关系,我们通过分析直接的递归拓扑排序来做到这一点,尽管效率不高,但可以遍历整个搜索空间。从本质上讲,拓扑排序算法是一个重复的过程,识别可用于调度的一组节点,并对该集合进行迭代以进行递归。在图论中,这组可用于调度的节点称为零入度集 ( z ),其中 ( z ) 是所有其入边和相应前驱节点(入度)均已调度的节点集合。图 5 演示了不同拓扑排序算法的递归树,其中树的高度是搜索步数,每从根到叶的路径都是一个拓扑排序 ( s ⊂ S_T )。该图突出了递归拓扑排序中的冗余 ( z ),然后将这些 ( z ) 合并为唯一的,识别为重复的签名,从而防止上述的重新计算。这使得对 ( z ) 的调度成为一个唯一的子问题,构成基于动态编程的拓扑排序。
图 5. 识别冗余的零入度集 z 并在拓扑排序算法中使 z 唯一方形,以减少重新计算
图 6. 可视化在搜索步骤 i = 8 中调度节点 u8 = H 。从 s8 、u8 和 upeak,8 开始,该图展示了算法如何计算 s9、u9 和 upeak,9