了解TensorFlow框架

2015年9月，Google于TensorFlow开源之际，发布了TensorFlow白皮书，介绍了TensorFlow的设计理念和实现方式。现在流行的大部分深度学习框架，都基于所谓的“数据流图”编程模型（又称“计算图”），为我们今后的编程提供了一种可选的编程范式。在本文中，我将以问答的形式解释关于TensorFlow的种种疑问。

Qustion 1：“数据流图”编程模型与传统编程模型的区别？

首先，“数据流图”的核心是一个有向图，图中的节点表示运算操作，边表示数据，整个图展现了数据的流动，因此称为数据流图。下图是一个示例。

在传统编程中，虽然我们也是对数据进行操作，但基本的三种控制逻辑“顺序、选择、循环”导致我们只能按照单一的流程处理数据，相当于数据流图只是一条线，而不是真正的图。换句话说，传统编程模型解决的是顺序操作流程，而数据流图则提供了并行计算的解决方案。其次，数据流图是数据驱动的，而不是指令驱动的。程序只规定数据的流向，而不能规定每一个操作何时执行，这就在另一个层面上提高了并行计算能力。最后，数据流图（我们这里只探讨静态数据流图）的定义和执行是分开的，用户不能像往常一样在某个操作处打断点查看输出内容，这削弱了该模型的调试能力，是为性能优化而付出的代价。

Qustion 2：TensorFlow的系统架构什么样？

TensorFlow框架包含三个模块，分别是client、master和worker，它们之间的逻辑关系如下图所示。

左侧为单机模式，右侧为分布式模式。client提供用户使用的编程接口，比如Python、C++ API等。master负责接受client的请求，构造数据流图，并分配任务给worker。在单机模式中，只有一个worker，负责计算数据流图中的所有计算任务。在分布式模式中，master需要为不同的worker分配不同的任务，以使总用时最小。

Question 3：分布式模式中，master如何合理地为多个worker分配任务？

master需要预先估算数据流图中各个节点的数据量和计算时长，估计数据量是为了保证各个设备的内存占用相差不大，估计计算时长是为了使总用时最短。这一步的算法在TensorFlow中称为Node Placement，即为每个节点寻找一个放置的位置。该算法会根据输入输出数据的规模、运算符种类来估算。此外，也可以根据实际运行过程中的实测结果来决定。

Question 4：分布式设备间的数据传输怎么实现？

当master为各个设备划分好任务后，这些设备间不可避免地要进行数据传输。TensorFlow的做法如下图所示。

首先，对跨机器的数据传输做了一层隔离，A设备中增加Send节点，用来对外发送数据，B设备中增加Receive节点，用来接收数据。这样，Send和Receive之间的通信与设备内部的通信可以使用完全不同的两套方案，简化了内部逻辑。实际实现中，Send和Receive之间通过TCP或RDMA的方式进行通信。

Question 5：如何scheduling？

当用户喂入数据，启动数据流图后，TensorFlow如何决定各个节点的执行顺序呢？朴素的想法是，每个节点执行完后通知与之相连的下一个节点，但如果下个节点有多个输入，它仍然无法启动，必须等到所有输入都到位后才能开始。因此TensorFlow使用了依赖计数的机制，每个节点记录其尚未满足的依赖的个数，当个数降为0时，启动该节点。这一方案完美体现了数据驱动的设计理念。

Question 6：反向传播的数据依赖怎么处理？

一旦考虑反向传播，数据流图就没那么简单了。虽然用户不必手动构造反向数据流，但TensorFlow会自动为我们构造，就像下图这个样子。

仔细观察上图，可以发现很多跨越数层的依赖，比如从MatMul到dAdd的灰色连线。这会导致GPU的内存占用急剧增大，因为几乎每个阶段计算的中间结果都不能丢弃。TensorFlow针对这种情况采取了若干种优化措施：(1)用更复杂的启发式算法调整图的执行顺序；(2)反向传播时重新计算前向传播的中间结果；(3)把中间结果保存到CPU内存中。

结语

Tensorflow开源至今，已经有了长足的发展。最新版本的实现可能与本文所述不再一致，但设计者的初衷不变，都是为了提供一个高扩展性、快速、高效的机器学习计算平台。

参考资料

TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems Google Research