OPTIMIZING DNN COMPUTATION WITH RELAXED GRAPH SUBSTITUTIONS

OPTIMIZING DNN COMPUTATION WITH RELAXED GRAPH SUBSTITUTIONS

现有的深度学习框架通过执行贪婪的基于规则的图变换来优化DNN模型的计算图，一般只考虑严格（strictly）提高运行时性能的变换。我们提出了relaxed graph substitutions，通过relaxing严格的性能提升约束，可以探索复杂的图优化，这大大增加了语义上等价的计算图的空间，可以通过重复应用一组合适的图变换来发现。我们引入了一种在一组relaxed graph substitutions上的回溯搜索算法来寻找优化网络，并使用基于流的图分割算法将计算图递归地分割成更小的子图，以实现高效搜索。我们在一个名为MetaFlow的系统中实现了relaxed graph substitutions，并表明MetaFlow比现有的深度学习框架分别提高了1.1-1.6x和1.1-1.2x的推理和训练性能。

1 INTRODUCTION

为了缓解日益增长的计算需求，在DNN中优化计算是标准的，它由数学运算符（如矩阵乘法、卷积等）的计算图定义。现有的深度学习系统如TensorFlow、PyTorch和TVM通过对图进行贪婪的基于规则的substitutions来优化输入计算图。

每一次替换都会用一个新的子图来替换与特定模式相匹配的子图，计算出相同的结果。例如，算子融合将几个算子合二为一，可以消除中间结果，从而减少系统开销，如内存访问和内核启动。现有的深度学习优化器考虑的是提高性能的替换，它们贪婪地反复应用于计算图，直到不能再进行替换。不考虑并非所有中间状态都是严格改进的更复杂的变换序列。因此，当前的优化器错过了许多更复杂的优化机会：我们表明，探索更大的变换空间可以比现有的基于规则的优化器提高广泛使用的DNN的性能高达1.6倍。

在本文中，我们提出了relaxed graph substitutions。我们通过relaxing the strict performance constraint来增加考虑的优化空间，允许任何保留语义的替换，无论它们是否提高性能。这些 "downgrading "的图替换作为变换图体系结构中间步骤，并最终发现运行时性能显著提高的新图，是非常有用的。为了有效地探索这个更大的计算图空间，我们在一组relaxed graph substitutions上使用回溯搜索，以发现多个替换步骤后的改进网络。作为一个例子，我们展示了如何使用我们的方法来优化广泛使用的ResNet架构（He等人，2016），如图1所示。最左边的图是贪婪的算子融合后的优化图，它将卷积和后续激活（即relu）结合成一个 “带激活的卷积”。然而，通过自适应地应用 relaxed graph substitutions （如图中箭头所示），可以在NVIDIA V100 GPU上生成比原始图形（最左边）快1.3倍的最终图形（最右边）。请注意，第一个图替换通过在kernel中填充额外的0，将卷积的核大小从1x1扩大到3x3。这降低了运行时的性能（因为具有较大内核的卷积运行速度较慢），但可以实现更多的后续内核融合，从而实现整体改进。第3节更详细地描述了其他图替换。

图1. ResNet模块上的relaxed graph substitutions。每个箭头都是一个图替换，相同颜色的虚线子图表示替换的源图和目标图。"conv axbxc "表示内核大小为a x b，输出通道为c的卷积。在NVIDIA V100 GPU上，最终图形（最右边）比原始图形（最左边）快1.3倍。

在现有DNN优化器中加入relaxed graph substitutions，并贪婪地应用它们，很容易导致性能下降。例如，图1中的放大算子substitutions很可能会降低性能，如果产生的卷积不能与另一个算子融合。虽然人们可以尝试通过在现有系统中添加特殊情况规则和启发式方法来解决这个问题，但我们认为这种方法在面对新的架构和new substitution rules时容易出错，而且很脆弱。取而代之的是，我们使用 cost-based的回溯搜索来有效探索计算图的大空间，而不需要优化器开发者实现大量新规则。

首先我们引入了一个cost model，该模型包含了多个成本维度（如FLOPs、执行时间、内存使用量等），可以准确估计不同计算图的性能。通过成本模型，我们可以快速比较不同的图形。

其次，我们提出了一种回溯搜索算法，可以快速找到小图的有效解。然而，SOTA的DNNs的计算图太大，无法直接探索所有等价计算图的搜索空间。因此，我们使用了一种graph split algorithm，将一个原始计算图递归地分割成多个较小的子图。 该图以最小化跨越不同子图的图替换数目的方式被分割，并且通过解决最大流问题来计算。这些子图通过回溯搜索进行优化，然后再拼接在一起，形成最终的优化图。图3描述了我们的图优化过程的概况。

我们在一个名为MetaFlow的系统中实现了relaxed graph substitutions，它可以用于优化任何现有深度学习框架的DNN计算图。特别是，我们表明TensorFlow、TensorFlow XLA和TensorRT可以直接使用MetaFlow的优化图来提高推理和训练性能。我们在五个real-world 的DNN上对MetaFlow进行了评估，包括Inception-v3（Szegedy等人，2016）、SqueezeNet（Iandola等人，2016）、ResNet-50（He等人，2016）、RNN文本分类（Kim，2014）和神经机器翻译（Wu等人，2016）。MetaFlow的搜索算法能够在5分钟内优化这些DNN。我们表明，MetaFlow优于现有的深度学习优化器，推理速度为1.1-1.6x，训练速度为1.1-1.2x。通过发现高效的计算图，使整体内存使用量减少了1.5倍，内核启动次数减少了3.3倍，从而实现了性能的提升。最后，我们表明MetaFlow的优化图可以直接输入到现有的框架中，并将其推理性能提高1.3倍。

总之，我们的贡献是：

我们引入了relaxed graph substitutions，它可以探索现有深度学习框架无法实现的复杂图优化。

我们提出了一种cost-based的搜索算法，可以在relaxed graph substitutions生成的搜索空间中自动找到优化的计算图。

我们实现了MetaFlow，第一个针对DNNs的 relaxed graph substitution优化器。在一个标准DNN的集合上，我们表明，与现有的框架相比，MetaFlow在保持相同网络精度的前提下，将运行时性能提高了1.1-1.6倍。

2 OVERVIEW

与现有的DNN优化器（Abadi等人，2016；Chen等人，2018；PyTorch）类似，MetaFlow使用计算图 $\mathcal{G}$ 来定义DNN模型中的计算和状态。每个节点是一个数学运算符（如矩阵乘法、卷积等），每个边是一个张量（即n维数组）。对于一个计算图 $\mathcal{G}$ ，取输入张量 $\mathcal{I}$ 并产生输出张量 $\mathcal{O}$，我们将其计算定义为 $\mathcal{O}=\mathcal{G}(\mathcal{I})$。我们定义两个计算图$\mathcal{G}$和${\mathcal{G}}^{\prime }$是等价的，如果$\mathcal{G}$和${\mathcal{G}}^{\prime }$对任意输入计算出数学上等价的输出的话 ($i.e.,\forall \mathcal{I}:\mathcal{G}(\mathcal{I})={\mathcal{G}}^{\prime }(\mathcal{I})$)。对于一个给定的计算图G，MetaFlow提供的图替换的合成，自动找到一个等价的计算图${\mathcal{G}}^{\prime }$，并优化运行时性能。

对于DNN模型来说，推理和训练程序是由不同的计算图定义的，如图2所示。推理图包括一个输入和一个或多个输出，而训练图一般有两个输入（即训练样本和标签）和多个输出（即每个算子中可训练参数的导数）。MetaFlow只是把推理和训练当作不同的图来优化，并在两个图上应用相同的技术。

图2. 一个4层示例CNN模型的推理和训练图。虚线是每个计算图的输入和输出。

图3显示了MetaFlow的主要组成部分。首先，对于任意输入的计算图，MetaFlow采用基于流的图分割算法，将输入图递归地划分为可供直接搜索的子图。

最后，MetaFlow以优化后的子图为basic building blocks，生成输入图的优化计算图。

MetaFlow是一个框架无关的计算图优化器：通过MetaFlow优化的计算图可以在各种深度学习运行时上执行，如TensorRT、TensorFlow和TensorFlow XLA。

3 RELAXED GRAPH SUBSTITUTIONS

本节介绍松弛图替换，每个图替换由一个source graph和一个target graph组成，source graph可以映射到DNN的计算图中的特定子图，target graph定义如何创建一个新的子图来替换映射的子图。

**Source graph. ** source graph定义替换的有效子图的结构。source graph中的每个节点都与一个类型相关联，并且只能被映射到同一类型的算子上。source graph还可以包括通配符节点，每个通配符节点可以映射到任何一个算子上。当一个算子的类型不影响the substitution procedure，通配符节点是有用的，它允许一个source graph描述多个相似的substitution方案。除了类型约束外，source graph还可以在一个或多个算子上加入额外的约束，以进一步限制映射。图4a演示了一个用于融合两个卷积的替换，它定义了对conv1和conv2的约束，以保证它们只能被映射到具有相同kernel size、stride和padding的卷积上。

图4. relaxed graph substitutions例子。(a)中的substitution用于图1的第二步（绿色），(b)中的substitution用于第三步（黄色）。

source graph中的边描述了算子之间的数据依赖关系。图替换要求映射的子图与source graph具有相同的数据依赖关系。每个算子可以有一个可选的external edge（如图4中的虚线边所示），它可以映射到连接到计算图中 external operators的零、一或多个边。

external edge表示算子的输出可以被external operators访问，并且必须在替换中保留。

Target graph. target graph描述了如何构造一个新的子图来代替映射的子图。对于每个新创建的算子，target graph定义如何通过使source graph中的参数和权重来设置参数和计算权重。对于source graph中的每个external edge，target graph中都有相应的external edge（也显示为dotted edges）。任何最初连接到source graph中映射算子的external operator现在都应连接到target graph中相应的算子。

Correctness. 如果一个图替换是有效的，那么它的source and target graphs对所有external edges计算出数学上相等的输出。这个定义类似于我们对等价计算图的定义，如果每个external edge都被视为图的输出。有效图替换的任何组合都保持生成的计算图之间的等价性。

**Composition. ** 许多复杂的图优化可以分解为一系列简单的relaxed graph substitutions。回想一下，图1展示了ResNet上的一个潜在优化，它通过扩大一个卷积的核来融合两个具有不同核大小的卷积。

图1. ResNet模块上的relaxed graph substitutions。每个箭头都是一个graph substitution，相同颜色的虚线子图表示替换的源图和目标图。"conv axbxc "表示内核大小为a x b，输出通道为c的卷积。在NVIDIA V100 GPU上，最终图形（最右边）比原始图形（最左边）快1.3倍。

作为另一个例子，下面的方程显示了如何通过使用一连串的graph substitutions来简化简单Simple Recurrent Unit的计算（Lei等人（2017）中的方程2和4），这些graph substitution\分别分配乘法，重新排序交换算子，并提出公共项。(by using a sequence of graph substitutions that distribute multiplications, reorder commutative operators, and factor out common terms, respectively.)

请注意，这两种优化都涉及复杂的图替换序列，需要在中间状态下暂时降低运行时性能。

4 THE METAFLOW SEARCH ALGORITHM

松弛图替换提供了一个潜在计算图的搜索空间，这些潜在计算图与初始计算图等价，但具有不同的运行时性能。在搜索空间中寻找最优图是一个挑战，因为搜索空间可能是无限的，这取决于所使用的替换规则。对于今天的DNN模型，穷尽地列举搜索空间肯定是不可行的。

本节描述了MetaFlow中使用的关键技术，这些技术可以有效地修剪搜索空间并快速找到优化的（但不一定是最优的）图。特别是，第4.1节引入了一个成本模型，该模型包含了多个成本维度（如FLOPs、执行时间、内存使用量等），并且能够准确预测各种计算图的执行性能。第4.2节介绍了一种回溯搜索算法，该算法能够在成本模型下的搜索空间中有效地找到一个优化的候选图。由于最新DNNs的计算图太大，无法直接优化，我们使用基于流的图分裂算法（第4.3节）递归地将计算图划分为更小的子图，同时最大化图替换机会。

4.1 Cost Model

#########

具体的cost model实现没介绍

#########

我们引入一个包含多个维度的成本模型来评估计算图的运行时性能。成本模型计算图中每个operator的度量(The cost model computes metrics for each operator in a graph)，并将其适当组合以获得总成本。这包括可以静态计算的度量（如FLOPs、内存使用量和内核启动次数）以及通常需要在特定硬件上测量的动态度量（例如，特定GPU或CPU上的执行时间）。

对于dynamic metrics，之前的工作表明，只需在硬件上测量几个有代表性的算子，就可以准确预测计算图的执行时间。由于大多数DNN算子都涉及无分支的稠密线性代数，因此在相同的条件下，它们在硬件上的性能是高度一致和可预测的参数。例如，一旦我们测量并存储了具有特定参数（即核大小、步长、padding等）的卷积的执行时间，我们就可以将该执行时间用于具有相同参数的其他卷积。

我们的cost model可以优化单个cost维度（例如，最小化overall FLOPs），也可以结合多个cost维度，例如在保持内存使用限制的情况下最小化执行时间（如果超过内存使用限制则返回infinite cost）。我们观察到，许多图替换导致了几个cost维度之间的权衡，而不是改善所有维度。例如，图4b中的减少了内存访问和内核启动，但代价是增加了FLOPs。

图4. relaxed graph substitutions例子。(a)中的substitution用于图1的第二步（绿色），(b)中的substitution用于第三步（黄色）。

4.2 Backtracking Search

我们现在描述一个回溯搜索算法，自动寻找成本模型下的优化计算图。算法1给出了伪码。所有候选图都被排队到一个全局优先级队列中，并按其代价的递增顺序进行出列。对于每个出列的图G，搜索算法通过对G应用潜在的图替换来生成新的图并将其入列。搜索算法使用参数$\alpha$（算法中的第13行）在搜索时间和最佳发现解之间进行折衷。通过设置$\alpha$=1，搜索算法成为一个简单的贪心算法，只考虑严格降低代价的图替换。随着$\alpha$增加，搜索算法会探索更大的搜索空间。

可以看到，所有的潜在图都排列成一个全局优先级队列，并且按照它们的成本以递增顺序移出队列。对于每个被移出队列的图 G，搜索算法都会通过在 G 上应用新的潜在图替代来生成和排列新的图。参数 α（算法 1 的第 13 行）被用于缩放当前最优 G 的损失。更近一步，我们可以直观地理解三种情况中的 α：

• (0< α <1)：这可以被认为是「鼓励」算法保持当前的解。只有当新图 G’ 的成本有明显改善时，算法才会更新当前最优的图。这有希望实现算法加速。
• (α = 1)：这与实现一个贪婪算法一样，且仅考虑严格降低成本的图替代。
• (α > 1) ：这可被视为「扩展」搜索空间，因为条件更加宽松。这个算法有希望找到最终能得到更优解的中间解。

回溯到思想出现在算法 1 的第 14 行。假设 α > 1，如果当前图 G’ 有更低的成本，这会被加回到队列中。这是因为未来的迭代有可能将 G’ 替换为另一个解 G’’，而 G’’ 有可能最终得到更差的解。因此，将 G’ 留在队列中能为算法提供追溯能力，从而防止当前路径不佳的问题。

4.3 Flow-Based Recursive Graph Split

############

图分割算法

############

许多SOTA的DNN模型太大，无法直接用回溯搜索进行优化。我们使用一个基于流的图分割算法来递归地将一个计算图分割成更小的不相交子图，这些子图适合回溯搜索。这是因为我们观察到图替换是在一些局部连通的算子上进行的，并且将计算图分解成更小的子图仍然可以保留大多数图替换。

为了将一个图分成两个不相交的子图，我们的目标是最小化两个子图之间的图替换数目，因为这些图替换不能在两个子图上执行。对于每个operator $o_i\in G$，我们将它的容量$Cap(o_i)$定义为映射到算子$o_i$的至少一个in-edge和一个outedge的图替换的数量。如果使用算子$o_i$拆分图，则禁用这些图形替换。通过使用$Cap(o_i)$作为每个算子的权重，我们将图分割问题映射为最小顶点切割问题(Cormen等，2009)，并且可以使用任何最大流量算法来寻找minimum cut。

最大流算法通过最小化生成图替换将任意图分解为两个不相交的子图。使用最大流算法作为子程序，算法2显示了一个图分割算法，该算法递归地将整个计算图分割为小于阈值的单个子图。

在运行回溯搜索算法来优化单个子图之后，MetaFlow将优化后的子图缝合在一起，构成一个完整的计算图。最后，围绕每个分割点执行局部回溯搜索，以获得跨越分割点的替换。

我们想指出的是，虽然flow-based的图分割算法是足够的，并且对于实验中使用的所有DNN都取得了良好的性能，但我们并不声称它是一种最优的图分割算法。我们研究了另一种图分割算法–平衡分割，看看结果是否不同。由于我们所研究的图中存在自然的分割点，这两种算法都实现了相同的性能。例如，在SqueezeNet中，我们的替换都没有越过火模块之间的边界(none of our substitutions cross the boundary between fire modules in SqueezeNet)，产生了一种简单的分割图的方法。然而，如果替换规则集或计算图不同，另一种图分割算法可能会被证明更有效。

5 IMPLEMENTATION

MetaFlow是一个针对任意计算图的框架无关的DNN优化器。MetaFlow的成本模型和运行时使用现有的深度学习库（例如，针对GPU的cuDNN（Chetlur等，2014）和cuBLAS（cuBLAS），以及针对CPU的MKL2）来估计计算图的执行时间，并在不同设备上进行实际执行。MetaFlow接受一个包含一个或多个成本维度的用户定义的成本函数，并找到一个优化成本函数的计算图。由MetaFlow优化的图可以自动转换为现有深度学习框架接受的格式，包括TensorRT、TensorFlow和TensorFlow XLA（TensorRT；Abadi等，2016）。这使得现有的深度学习框架可以直接使用MetaFlow的优化图作为输入，以提高运行时性能。将现有的深度学习框架的运行时性能提高了1.3倍，尽管这些系统在执行输入计算图之前，会在内部进行基于规则的图转换。