AI 前线导读: 一些最新机器学习模型, 特别是 RNN 和强化学习, 其训练和推理中使用了细粒度的动态控制流. 在这些模型中, 循环体执行和依赖于数据的条件执行需要动态控制流的支持, 实现在分布异构环境各计算设备 (包括 CPU,GPU, 以及 TPU 等定制的 ASIC) 间作快速的控制流决策. Google Brains 团队认为, 出于性能, 可扩展性和运算表达能力上的考虑, 现代机器学习系统应该支持分布式异构计算环境中的动态控制流.
本文是一篇 EuroSys 2018 会议接收论文, 论文的主要贡献者为 Yuan Yu(本文工作主要完成于在 Google Brain 工作期间), 以及 Google Brain 团队 Jeff Dean 等人, 该团队实现并开源了 TensorFlow 系统.
论文提出了一种支持动态控制流的分布式机器学习编程模型, 并介绍了该模型的设计及在 TensorFlow 中的实现. 为表示机器学习运算, 该编程模型对数据流图 (data graph) 做了扩展, 实现了多个显著的特性. 第一, 该编程模型支持对数据流图的细粒度分区, 实现了在多个异构设备上运行条件分支和循环体; 第二, 使用该模型编写的程序支持自动微分 (automatic differentiation) 计算和分布式梯度计算, 这些计算是训练机器学习模型训练所必须的; 第三, 模型采用了非严格语义(non-strict semantics), 支持多个循环迭代在设备间的并行执行, 实现了计算和 I/O 操作的交叠(overlapping).
Synopsys 的唐杉博士点评道:"传统神经网络大部分是静态图, 这也是很多专用加速器能够高效加速神经网络的一个基本前提. 如果是一个动态的控制流, 则底层架构设计可能有巨大差别, 运算, 存储和通信这三要素在硬件中的比例又会发生新的变化." 本文提出的动态控制流编程模型, 将给 AI 硬件设计带来新的挑战.
来源: https://juejin.im/post/5b0cc6e751882515731af6c5