【AI系统】计算图的控制流实现

计算图的控制流实现

计算图在数学上作为一个有向无环图（DAG，Directed Acyclic Graph），能够把神经网络模型的概念抽象出来作为同一描述，不过在计算机的编程中，会遇到很多 if、else、while、for 等控制流语句，有向无环图改如何表示控制流变成了计算图中一个很重要的问题。好处在于，引入控制流之后，开发者可以向计算图中引入分支选择以及循环控制逻辑，进而构造出更加复杂的神经网络模型结构。

目前，以 Pytorch 为例，它支持的仅仅是 Python Control Flow，即在 Python 层执行控制逻辑，而非计算图中支持控制流。这样就存在一个问题，如果要部署带 Control Flow 的模型就会比较困难，如何灵活部署带控制流的计算图到不支持 Python 的设备上？

计算图中的控制流实现，与控制流图并不是一个概念。在计算机科学中，控制流图 (CFG) 是程序执行期间所有可能路径的图形表示。控制流图是由 Frances E. Allen 发现的。他指出，Reese T. Prosser 之前曾使用布尔连接矩阵进行流分析。CFG 是许多编译器优化和静态分析工具不可或缺的一部分。

控制流

AI 框架作为一个可编程系统，在设计时一个首要设计选择是如何让开发者，能够独立于实现细节以最自然的方式描述出各类神经网络模型。描述的完备性不仅影响 AI 框架能所够支持的神经网络结构，决定了开发者在使用高级编程语言，去实现神经网络模型的灵活性，也影响 AI 框架后端优化的技术手段。

背景

在计算机科学中，控制流（Control Flow）定义了独立语句，指令，函数调用等执行或者求值的顺序。例如，根据函数 A 的输出值选择运行函数 B 或者 C 中的一个。

AI 框架把神经网络的计算过程，抽象为有向无环图。使用有向无环图描述神经网络计算的方式，符合算法开发者对神经网络的概念定义：算子间拓扑结构对学习特性有重要影响，可以通过计算图，方便地描述出大多数通过堆叠深度或多分枝形成的复杂神经网络。然而，随着神经网络算法的快速发展，一些新颖的神经网络结构很难自然地表示为纯计算图。

以 Transformer 结构的神经网络为例，来看看使用最自然地方式描述这些算法对 AI 框架会带来什么新的表示要求。Transformer 结构的神经网络算法中，图的左侧是一个通用 Transformer 结构的中关键步骤的计算示意图，通过堆叠 Transformer 结构使得网络模型层数更深，右侧对应了使用最自然的方式描述这一算法计算过程的伪代码。

Transformer 是一种基于注意力机制的神经网络结构，由多个 Encoder 和 Decoder 堆叠而成，可以应用于各种自然语言处理任务。在具体应用时，可以根据任务的特点和需求，选择不同的 Transformer 架构来构建模型。从伪代码描述中可以看到，想要以通用的方式，自然地描述出 Transformer 的算法框架，均依赖于循环控制逻辑 for。

难点

引入控制流将会使得计算图的构建以及前向传播带来很大的差异。

首先，计算图将变为动态的方式，分支选择以及循环控制流只有在真实运行的时候，才能够依据其依赖的数据输入来判断走哪个分支、是否结束循环。

其次，控制流引入的另一个难点在于循环控制流的实现。引入循环之后，原本的计算图在逻辑上出现了环，从而无法进行有效的拓扑排序。所以对于有控制流的计算图，前向计算和反向传播的实现要么抛弃拓扑排序这一思路，要么就要通过其他手段将循环进行拆解。

为了能够支持含有控制流结构的神经网络模型，AI 框架需要引入了对动态控制流这一语言结构（language construct）的支持。目前基于计算图，在控制流解决方案上，主要采用了三类设计思路：

• 复用宿主语言：复用前端宿主语言的控制流语言结构，用前端语言中的控制逻辑驱动后端计算图的执行；

• 支持控制流原语：AI 框架的后端对控制流语言结构进行原生支持，计算图中允许计算流和控制流混合；

• 源码解析：前端对高级编程语言的代码进行解析称计算图，后端对控制流语言结构解析成子图，对计算图进行延伸。

复用宿主语言以 PyTorch 为典型代表，支持控制流原语以 TensorFlow 为典型代表，源码解析的方式则以 MindSpore 为典型代表。

动态图

每一次执行神经网络模型，AI 框架会依据前端编程语言描述，动态生成一份临时的计算图（实际为单算子），这意味着该实现方式下计算图是动态生成，并且过程灵活可变，该特性有助于在神经网络结构调整阶段提高效率，这种实现方式也被称为动态计算图。