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
网友评论