本节是自动求导框架技术的第三节,本系列其余文章包括
1. 控制流简介
前面两节我们已经知道了如何实现普通计算图,然后基于计算图的拓扑排序和逆向拓扑排序完成前向传播与反向求导。但是仅有这些是不够的,在 tensorflow 中还引入了一种叫做控制流的机制来提高框架的灵活性。引入控制流之后,框架的使用者可以向计算图中引入分支选择以及循环,进而构造出更加复杂的神经网络结构。
2. 控制流的实现难点
引入控制流将会使得计算图的构建以及前向传播带来很大的不同。首先,计算图将变得动态。分支选择以及循环控制流只有在真实运行的时候依据其依赖的数据的输入来判断走哪个分支、是否结束循环。就好比建筑师得到了一张动态的图纸,在建高楼之前完全不知道大楼会盖到多少层,长什么样子;只有真正建到某些层之后考虑了一下现在建筑的状况,然后决定接下来的部分怎么建造。
控制流引入的另一个难点在于循环控制流的实现。引入循环之后,原本的计算图在逻辑上出现了环,从而无法进行拓扑排序。所以对于有控制流的计算图,前向计算和反向传播的实现要么抛弃拓扑排序这一思路,要么就要通过某些方式将循环进行拆解。
3. 控制流的实现思路
3.1 虚拟图的结构
为了解决计算图动态的问题,我引入了“虚拟图”这一概念。框架的使用者并不直接构造计算图,而是构造一个包含了控制流的虚拟图,然后框架对虚拟图进行前向传播的同时动态生成一个计算图,之后基于生成的计算图进行反向传播。虚拟图的引入解决了动态计算图难以进行反向传播的问题。
虚拟图的组成包含三种节点:分支结点 branch_node,循环节点 loop_node,普通节点 virtual_node。branch_node 可以决定向其子节点提供哪个计算节点的依赖(包括空依赖)从而实现分支控制,接下来主要谈一下 loop_node 的问题。对于虚拟图我依然希望其是有向无环图,所以我需要将逻辑上的循环和有向无环图这两个矛盾进行调和,实现方法的思路来自于图的强连通分量。参考算法导论中文版第22章的引理22.13。
对于一个图而言,每个强连通分量可以看做是一个广义的节点,这些广义节点构成的图是一个有向无环图。
所以解决虚拟图包含循环的方式就是把 loop_node 看做一个广义节点,其中包含了一个“子虚拟图”,这个子虚拟图在 loop_node 的中的 inner_loop () 函数的控制下可以循环遍历多次,直到满足 loop_node 的循环跳出条件。由于 loop_node 中包含的子虚拟图也是一个虚拟图,所以整个虚拟图是一个递归层层嵌套的结构,可以实现大部分的循环级联、循环嵌套。但是虚拟图在其各个层面上依然是有向无环图,所以可以进行拓扑排序。
3.2 虚拟图如何生成计算图
虚拟图中的 build_compute_graph (cg, idx) 函数用于构造计算图,cg是待生成计算图的引用,表示 compute_graph;idx 是循环的下标。
在虚拟图中,每个节点用一个二元组唯一标识:<type, id>,其中 type 标识这个虚拟节点的类别,根据type 虚拟节点能够知晓自己应该生成什么样的计算节点;id 是用于区分同一个类别的不同虚拟节点。由于虚拟图中存在循环,一个虚拟节点可能在某个循环中生成多个计算节点,所以计算节点使用一个三元组唯一标识:<type, id, idx>,idx 就是上面提到的循环下标。
虚拟图中的 build_compute_graph (cg, idx) 流程介绍如下:
1. 对虚拟图拓扑排序;
2. 对于拓扑排序后的每个节点:
if (是loop_node) : 执行 这个 loop_node 的 inner_loop (cg) 函数
if (是普通virtual_node):首先查找当前待生成的计算节点的依赖节点是否齐全(由于分支结点的原因当前虚拟节点的路径可能并不会被经过),如果依赖节点齐全,则根据虚拟节点的 type 值生成相应的计算节点,加入计算图 cg,并调用计算节点中的 op () 方法完成这个计算节点输出的计算。
3. 返回该计算图的末尾节点。
对于loop_node 的 inner_loop (cg) 函数,其主要流程如下:
1. 调用 loop_node 中用户自定义的 init (cg) 函数完成循环的初始化;
2. while (用户自定义condition (cg) 函数返回0) {
循环下标 idx 自增
遍历 loop_node 中的子虚拟图,即调用子虚拟图的 build_compute_graph (cg, idx) 函数
}
3. 记录下循环的结果——最后生成的计算节点
通过对虚拟图前向传播的同时构建出计算图,这个计算图是一个有向无环图。然后反向传播的就可以依据该计算图执行。