对TensorFlow中数据流图的可视分析

TensorFlow是谷歌开发的、当下最流行的机器学习软件库之一。它采用数据流图（Dataflow Graph）来表达机器学习算法的计算过程，用户可以定制不同的数据流图来构建自己的算法。然而，随着深度学习的兴起与流行，各类神经网络渐趋大规模、复杂化。算法开发者仅凭借自身的理解与记忆、很难把握算法的各部分体系结构，相互之间也难以进行沟通。为此，这篇文章[1]提出了可视化工具TensorFlow Graph Visualizer，通过可视分析帮助用户在TensorFlow中进行算法分析与开发。值得一提的是，该文章荣获了IEEE VAST 2017的最佳论文奖（Best Paper Award）。

1. TensorFlow中的数据流图

数据流图是TensorFlow中最重要也最基础的概念，它将数据看作“流”、将数据处理算法看作节点，从而任意一种复杂算法都能够用一张图来表示。而TensorFlow使用张量（Tensor）来表示所有类型的数据，相应的数据流图便称为Tensor Flow Graph。

图1 TensorFlow中数据流图的点边类型

在Tensor Flow Graph中，各有三种类型的节点和边（如图1）。其中三种节点包括：数据操作（Operation）、数据赋值（Initializer）、以及日志记录（Log）。三种连边包括：数据值（Value）、数据索引（Reference）、以及算法依赖关系（Control，如A操作必须在B操作之前）。

数据流图是TensorFlow中最重要也最基础的概念，它将数据看作“流”、将数据处理算法看作节点，从而任意一种复杂算法都能够用一张图来表示。而TensorFlow使用张量（Tensor）来表示所有类型的数据，相应的数据流图便称为Tensor Flow Graph。

2. 力导向图布局

在进一步设计之前，作者利用力导向布局算法对数据流图进行了直接的可视化，但效果并不理想（如图2）。

图2 数据流图的力导向布局

其中图的区块化结构不明显，连边之间存在许多交错，而且布局极不稳定、局部的节点更新会扰动整个布局。究其原因，在于以下两点：

1) 图中存在大量负责日志记录的“中介节点”，将许多本不相关的子图紧密联系在一起。具体来说，每个子算法都有日志记录模块（图2中蓝色节点），它们与唯一的日志书写模块（图2中红色节点）相连。这些连边并不具备算法分析上的价值，却会破坏各个算法子图之间的独立性，需要在后续优化中去除。

2) 图布局缺乏层次结构。由于所有节点都被同等对待，对任意节点的更新都免不了影响所有其他节点，导致布局不稳定。要解决这个问题，需要为图建立起层次结构（Hierarchy），通过“全局定位+局部更新”的方式来避免布局的频繁扰动。此外，建立层次也能使算法的逻辑结构变得更为清晰。

3. Tensor Flow Graph Visualizer

针对直接图布局的各个缺点，文章作者提出了Tensor Flow Graph Visualizer，针对TensorFlow中的数据流图进行了结构与布局上的优化。优化分为三步：

1) 去除冗余节点

正如第2节所述，数据流图中存在大量负责日志记录的“中介节点”、扰乱了原本各自独立的算法子图。这些日志记录节点并不具备分析意义，我们可以将它们从图中去掉。此外，常量节点只输出到其他操作而没有任何输入，结构非常简单，同样可以被去掉。移除冗余节点后，力导向布局得到了很大的优化（如图3），其中的子图结构已经能够相互分离并显现出来。

图3 去除冗余节点后的数据流图

2) 建立层次结构

表达算法的数据流图，其本身是存在层次化逻辑结构的，即多个元操作（Meta Operation）能够组合成子算法/子步骤、并自底向上地构成一个更为复杂的算法。在神经网络中，各个中间层便是操作的组合，其中的元操作，就是各神经元中的数据加权和激活函数的计算。

图4 算法操作的层次化显示

层次化组织后的数据流图只显示顶层的命名空间，隐藏了其中的底层操作，并将数据流相应地捆绑起来（如图4(a)）。其中椭圆形节点则表示元操作，圆角矩形表示一组操作，矩形高度则表示组内操作的数目。算法自动检测了内在结构相同的子图，并使用同一种颜色进行标记，灰色则表示内部结构不同。连边中，虚线表示算法依赖关系，实线表示数据传输，连边粗细表示数据流的大小。用户可以双击打开一组操作、以观察其内部的操作细节（如图4(b)）。在用户展开一组操作、或更新其中某些节点时，只有组内的子图布局会发生变化，保证了图布局的稳定。

建立层次结构最直接的方法是对节点进行聚类，但聚类的结果可能不符合开发者的期望。为此，作者选择了从代码本身提取层次结构，即利用各个操作的命名空间（namespace）来对节点进行分组。相应地，开发者需要自行决定代码的层次结构，并将其反应在操作的命名空间中。

3) 提取辅助节点

图5 提取辅助节点以简化图结构

在前两步优化以后，图中依然存在部分度数较高的节点，容易造成不同子图之间的相互耦合（如图5(a)）。这些节点多是变量赋值、数据记录等辅助操作，虽然并非至关重要，但也不可或缺、无法直接去除。

为此，作者选择提取这些辅助性节点、将其画在流图之外，并通过克隆节点的小图标（proxy）来表现节点之间的相连关系（如图5(b)）。提取了度数高的辅助节点后，不同子图能够进一步地解耦合，大为简化了数据流图的整体结构。

4. 案例分析

用Tensor Flow Graph Visualizer展现卷积神经网络（Convolutional Neural Network, CNN）LeNet-5模型的算法结构，如图6所示。

图6 CNN算法模型的数据流图

自下而上可以依次看见图片的读取（read）、预处理（process_image）、批处理（shuffle_batch）、两个结构相同（具有相同颜色）的卷积层（conv1 & conv2）、两个结构不同的全连接层（fully3 & fully4）、作为输出的激活函数层（softmax_linear）、以及计算反馈的残差操作（cross_entropy）等等。其中同样可以看见如初始化（init）、池化（pooling）、归一化（norm）等元操作。可见，利用文章所提出的图构建和布局方法，能较好地从源代码中提取出深度学习算法的整体逻辑结构。用户通过交互探索，也能够接触到底层的操作细节，方便开发者对算法进行更细致的调试和检错。

5. 总结

这篇文章针对TensorFlow中纷繁复杂的数据流图，从不同方面进行了优化，并向开发者提供了一个结构清晰、细节充分的层次化算法视图，大大提高了用户在TensorFlow中设计、实现、调试深度学习算法的效率。不同于其他基于神经网络结构、或是训练优化过程的可视化方法，该方法借助数据流图强大的表达能力，能够从顶层概念到实现细节等多个层面来表现深度学习算法，其思路值得我们借鉴和学习。

文献：

[1] Wongsuphasawat, K., Smilkov, D., Wexler, J., Wilson, J., Mané, D., Fritz, D., ... & Wattenberg, M. (2018). Visualizing Dataflow Graphs of Deep Learning Models in TensorFlow. IEEE transactions on visualization and computer graphics, 24(1), 1-12.

北京大学可视化/可视分析博客(http://vis.pku.edu.cn/blog/)

我们将在这里和大家共享对可视化与可视分析研究领域的理解，一同探讨可视化研究的进展，观察和评论最新的可视化技术和应用。我们关心可视化与可视分析核心的成就，也关注相关领域的发展。欢迎关注！