TF-Replicator:研究人员的分布式机器学习

时间:2019-03-10 14:05:32 阅读:11次

  在DeepMind,研究平台团队构建基础架构以增强和加速我们的AI研究。今天,我们很高兴分享我们如何开发TF-Replicator,这是一个软件库,可帮助研究人员以最小的努力和以前没有分布式系统的经验,在GPU和云TPU上部署他们的TensorFlow模型。 TF-Replicator的编程模型现已作为TensorFlow的tf.distribute.Strategy的一部分开源。此博客文章概述了TF-Replicator背后的想法和技术挑战。有关更全面的描述,请阅读我们的arXiv论文。

  近期AI突破的一个反复出现的主题 - 从AlphaFold到BigGAN再到AlphaStar--需要轻松可靠的可扩展性。越来越多的计算能力使研究人员能够利用新功能培养更大的神经网络。为了解决这个问题,研究平台团队开发了TF-Replicator,它允许研究人员针对机器学习定位不同的硬件加速器,将工作负载扩展到许多设备,并在不同类型的加速器之间无缝切换。虽然它最初是作为TensorFlow上的库开发的,但TF-Replicator的API已经集成到TensorFlow 2.0的新tf.distribute.Strategy中。

  虽然TensorFlow为CPU,GPU和TPU(Tensor Processing Unit)设备提供直接支持,但在目标之间切换需要用户付出巨大努力。这通常涉及为特定硬件目标专门化代码,将研究思想限制在该平台的功能上。一些现有的框架构建在TensorFlow之上,例如估算师,寻求解决这个问题。但是,它们通常针对生产用例,缺乏快速迭代研究思路所需的表现力和灵活性。

构建分布式机器学习库

  我们开发TF-Replicator的最初动机是为DeepMind研究人员提供一个使用TPU的简单API。 TPU为机器学习工作负载提供可扩展性,通过我们的BigGAN模型实现研究突破,例如最先进的图像合成。 TensorFlow的TPU原生API与GPU的目标不同,形成了TPU采用的障碍。 TF-Replicator提供了一个更简单,更用户友好的API,隐藏了TensorFlow TPU API的复杂性。重要的是,研究平台团队与各种机器学习领域的研究人员密切合作开发了TF-Replicator API,以确保必要的灵活性和易用性。

TF-Replicator API

  使用TF-Replicator编写的代码与TensorFlow中为单个设备编写的代码类似,允许用户自由定义自己的模型运行循环。用户只需要定义(1)公开数据集的输入函数,以及(2)定义模型逻辑的步骤函数(例如梯度下降的单步):

#使用repl.context()部署模型,使用TpuReplicator.repl = tf_replicator.TpuReplicator(num_workers = 1,num_tpu_cores_per_worker = 8):model = resnet_model()base_optimizer = tf。 train.AdamOptimizer()optimizer = repl.wrap_optimizer(base_optimizer)#...用tf.train.MonitoredSession定义副本input_fn和step_fn.per_replica_loss = repl.run(step_fn,input_fn)train_op = tf.reduce_mean(per_replica_loss)的代码)作为会话:在xrange中的i的repl.init(会话)(num_train_steps):session.run(train_op)repl.shutdown(会话)

  将计算扩展到多个设备需要设备彼此通信。在训练机器学习模型的背景下,最常见的通信形式是累积梯度以用于优化算法,例如随机梯度下降。因此,我们提供了一种方便的方法来包装TensorFlow优化器,以便在更新模型的参数之前在设备之间累积渐变。对于更一般的通信模式,我们提供类似MPI的原语,例如`all_reduce`和`broadcast`。这些使得实现全局批量标准化等操作变得微不足道,这种技术对于扩展BigGAN模型的培训至关重要(参见本文第3节)。

全屏 fullscreen_mobile 输入数据从主机发送到每个GPU,立即开始处理。当需要在GPU之间交换信息时,它们会在发送数据之前进行同步。

实现

  对于多GPU计算,TF-Replicator依赖于“图中”复制“模式,其中每个设备的计算在同一TensorFlow图中复制。通过连接来自设备的相应子图的节点来实现设备之间的通信。在TF-Replicator中实现这一点具有挑战性,因为通信可以在数据流图中的任何点进行。因此,计算构造的顺序至关重要。

  我们的第一个想法是在一个单独的Python线程中同时构建每个设备的子图。遇到通信原语时,线程同步,主线程插入所需的跨设备计算。之后,每个线程将继续构建其设备的计算。但是,当我们考虑这种方法时,TensorFlow的图形构建API不是线程安全的,这使得在不同线程中同时构建子图非常困难。相反,我们使用图形重写来在构建所有设备的子图之后插入通信。构建子图时,占位符将插入需要通信的位置。然后,我们跨设备收集所有匹配的占位符,并用适当的跨设备计算替换它们。

全屏 fullscreen_mobile 当TF-Replicator构建图中复制计算时,它首先独立地为每个设备构建计算,并留下占位符,用户指定了跨设备计算。一旦构建了所有设备的子图,TF-Replicator就会通过用实际的跨设备计算替换占位符来连接它们。

在DeepMind上构建AI研究平台

  通过在TF-Replicator的设计和实施过程中与研究人员密切合作,我们能够构建一个库,允许用户轻松地在多个硬件加速器上进行计算,同时使他们具备切割所需的控制和灵活性 - 人工智能研究。例如,我们在与研究人员讨论后添加了MPI风格的通信原语,例如all-reduce。 TF-Replicator和其他共享基础架构使我们能够在强大的基础上构建越来越复杂的实验,并在整个DeepMind中快速传播最佳实践。

  在撰写本文时,TF-Replicator是DeepMind中使用最广泛的TPU编程接口。虽然图书馆本身并不局限于训练神经网络,但它最常用于大批量数据的训练。例如,BigGAN模型在尺寸为2048的批次上训练,最多可达到TPUv3吊舱的512个核心。在具有分布式演员 - 学习者设置的强化学习代理中,例如我们的重要性加权的演员 - 学习者架构,通过让许多演员通过与环境交互来产生新体验来实现可伸缩性。然后,学习者处理该数据以改进代理的策略,表示为神经网络。为了应对越来越多的参与者,TF-Replicator可以用来轻松地将学习者分配到许多硬件加速器上。这些和其他示例在我们的arXiv论文中有更详细的描述。

   TF-Replicator只是DeepMind研究平台团队构建的众多有影响力技术的例子之一。从AlphaGo到AlphaStar,DeepMind在人工智能方面的许多突破都是由团队实现的。如果您分享我们的使命并对加速最先进的人工智能研究感到兴奋,请通过https://deepmind.com/careers 了解研究平台中的开放式软件工程职位(机器学习经验是可选的这些角色)。

   这项工作由DeepMind的研究平台团队完成。我们要感谢Frederic Besse,Fabio Viola,John Aslanides,Andy Brock,Aidan Clark,SergioGómezColmenarejo,Karen Simonyan,Sander Dieleman,Lasse Espeholt,Akihiro Matsukawa,Tim Harley,Jean-Baptiste Lespiau,Koray Kavukcuoglu,Dan Belov和DeepMind的许多其他人在TF-Replicator的整个开发过程中提供了宝贵的反馈。我们还要感谢Priya Gupta,Jonathan Hseu,Josh Levenberg,Martin Wicke和Google的其他人,他们将这些想法作为tf.distribute.Strategy的一部分提供给所有TensorFlow用户。

TF-Replicator:研究人员的分布式机器学习所属专题:机器学习 《TF-Replicator:研究人员的分布式机器学习》链接:http://www.zhouchuanxiong.net/1065

TF-Replicator:研究人员的分布式机器学习相关文章