DeepMind又放福利：开源了一个内部的分布式机器学习库TF-Replicator
今天，deepmind又放福利：开源了一个内部的分布式机器学习库tf-replicator，可以帮助研究人员将tensorflow模型轻松部署到gpu、tpu，并实现不同类型加速器之间的无缝切换。
最近ai领域的突破，从alphafold到biggan再到alphastar，一个反复出现的主题是，对方便、可靠的可扩展性的需求。
研究人员已经能够获取越来越多计算能力，得以训练更大的神经网络，然而，将模型扩展到多个设备并不是一件容易的事情。
今天，deepmind又将其内部一个秘密武器公之于众——tf-replicator，一个可以帮助研究人员将他们的tensorflow模型轻松部署到gpu、cloud tpu的分布式机器学习框架，即使他们之前完全没有使用分布式系统的经验。
tf-replicator由deepmind的研究平台团队开发，初衷是为deepmind的研究人员提供一个简单的接入tpu的api，现在，tf-replicator已经是deepmind内部最广泛使用的tpu编程接口。
tf-replicator允许研究人员针对机器学习定位不同的硬件加速器进行，将工作负载扩展到许多设备，并在不同类型的加速器之间无缝切换。
虽然它最初是作为tensorflow上面的一个库开发的，但目前tf-replicator的api已经集成到tensorflow 2.0新的tf.distribute.strategy中，作为 tf.distribute.strategy的一部分开源：
https://www.tensorflow.org/alpha/guide/distribute_strategy
团队还公开了相关论文：tf-replicator: distributed machine learning for researchers，全面描述了这个新框架的技术细节。
https://arxiv.org/abs/1902.00465
接下来，我们将介绍tf-replicator背后的想法和技术挑战。
构建一个分布式机器学习库
虽然tensorflow为cpu、gpu和tpu设备都提供了直接支持，但是在目标之间切换需要用户付出大量的努力。这通常涉及为特定的硬件目标专门编写代码，将研究想法限制在该平台的功能上。
一些构建在tensorflow之上的现有框架，例如estimators，已经试图解决这个问题。然而，它们通常针对生产用例，缺乏快速迭代研究思路所需的表达性和灵活性。
我们开发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)一个定义其模型逻辑的step函数(例如，梯度下降的单个step):
# deploying a model with tpureplicator.repl = tf_replicator.tpureplicator( num_workers=1, num_tpu_cores_per_worker=8)with repl.context(): model = resnet_model() base_optimizer = tf.train.adamoptimizer() optimizer = repl.wrap_optimizer(base_optimizer)# ... code to define replica input_fn and step_fn.per_replica_loss = repl.run(step_fn, input_fn)train_op = tf.reduce_mean(per_replica_loss)with tf.train.monitoredsession() as session: repl.init(session) for i in xrange(num_train_steps): session.run(train_op) repl.shutdown(session)
将计算扩展到多个设备需要设备之间进行通信。在训练机器学习模型的背景下，最常见的通信形式是累积梯度(accumulate gradients)以用于优化算法，如随机梯度下降。
因此，我们提供了一种方便的方法来封装tensorflow optimizers，以便在更新模型参数之前在设备之间累积梯度。对于更一般的通信模式，我们提供了类似于mpi的原语，如“all_reduce”和“broadcast”。这些使得实现诸如全局批标准化之类的操作变得非常简单，这是扩展biggan模型训练的关键技术。
输入数据从主机发送到各个gpu, gpu立即开始处理。当需要在gpu之间交换信息时，它们会在发送数据之前进行同步。
实现
对于多gpu计算，tf-replicator依赖于“图内复制”(“in-graph replication)模式，其中每个设备的计算在同一个tensorflow graph中复制。设备之间的通信是通过连接设备对应子图中的节点来实现的。在tf-replicator中实现这一点很具挑战性，因为在tensorflow graph中的任何位置都可能发生通信。因此，构造计算的顺序至关重要。
我们的第一个想法是在一个单独的python线程中同时构建每个设备的子图。当遇到通信原语时，线程同步，主线程插入所需的跨设备计算。之后，每个线程将继续构建其设备的计算。
然而，在我们考虑这种方法时，tensorflow的图形构建api不是线程安全的，这使得在不同线程中同时构建子图非常困难。相反，我们使用图形重写(graph rewriting)在所有设备的子图构建完成后插入通信。在构造子图时，占位符被插入到需要通信的位置。然后，我们跨设备收集所有匹配占位符，并用适当的跨设备计算替换它们。
当tf-replicator构建一个in-graph replicated计算时，它首先独立地为每个设备构建计算，并将占位符留给用户指定的跨设备计算。构建好所有设备的子图之后，tf-replicator通过用实际的跨设备计算替换占位符来连接它们。
为ai研究构建一个平台
通过在tf-replicator的设计和实现过程中与研究人员密切合作，我们最终构建一个库，让用户能够轻松地跨多个硬件加速器进行大规模计算，同时让他们拥有进行前沿ai研究所需的控制和灵活性。
例如，在与研究人员讨论之后，我们添加了mpi风格的通信原语，如all-reduce。tf-replicator和其他共享基础架构使我们能够在稳健的基础上构建越来越复杂的实验，并在整个deepmind快速传播最佳实践。
在撰写本文时，tf-replicator已经成为deepmind应用最广泛的tpu编程接口。虽然这个库本身并不局限于训练神经网络，但它最常用来训练大量数据。例如，biggan模型是在一个512核的tpuv3 pod训练的，batch size为2048。
在采用分布式actor-learner设置的增强学习智能体中，例如我们的重要性加权actor-learner架构，可扩展性是通过让许多actor通过与环境的交互生成新的体验来实现的。然后，learner对这些数据进行处理，以改进agent的策略，表示为一个神经网络。为了应对越来越多的actor，tf-replicator可以很轻松地将learner分布在多个硬件加速器上。
这些以及更多例子在我们的arxiv论文中有更详细的描述。
blog:
https://deepmind.com/blog/tf-replicator-distributed-machine-learning/
paper：
https://arxiv.org/abs/1902.00465
2019世界智能网联汽车大会将引发汽车新时代的潮流
电压跟随器的作用
抢单台积电,三星晶圆代工的底气何在?
中国最大的工程机械公司徐工，和中国最大的空调企业格力都遭遇过出售的生死危机
大话西游手游联手冰峰，国风老字号重焕江湖新情义
DeepMind又放福利：开源了一个内部的分布式机器学习库TF-Replicator
小米PK华为：智能家居这场战争谁能笑到最后
拒绝等待，华为云CDN下载加速就是要快人一步
99元的Xisem西圣ASN蓝牙耳机，真的香！
0xUniverse : 除了脚下的土地，我们还拥有广阔的星空
51单片机的汇编程序库
矿用联网机锚网气动联扣机CQ24CQ50C4530型号齐全有煤安
爆料称明年新款iPhone刘海将缩小
广和通与多家企业联合成立深圳市5G-A应用创新产业联盟
功率放大器的性能特点和使用步骤
地平线与智能汽车计算平台方案商国汽智控签署战略合作协议
LED灯珠内氯的来源及解决办法
典型案例中电源管理解决方案应用探讨
为什么一定是三次握手和四次挥手
3月新能源汽车上险数据的一些基本分析
最近ai领域的突破，从alphafold到biggan再到alphastar，一个反复出现的主题是，对方便、可靠的可扩展性的需求。
研究人员已经能够获取越来越多计算能力，得以训练更大的神经网络，然而，将模型扩展到多个设备并不是一件容易的事情。
今天，deepmind又将其内部一个秘密武器公之于众——tf-replicator，一个可以帮助研究人员将他们的tensorflow模型轻松部署到gpu、cloud tpu的分布式机器学习框架，即使他们之前完全没有使用分布式系统的经验。
tf-replicator由deepmind的研究平台团队开发，初衷是为deepmind的研究人员提供一个简单的接入tpu的api，现在，tf-replicator已经是deepmind内部最广泛使用的tpu编程接口。
tf-replicator允许研究人员针对机器学习定位不同的硬件加速器进行，将工作负载扩展到许多设备，并在不同类型的加速器之间无缝切换。
虽然它最初是作为tensorflow上面的一个库开发的，但目前tf-replicator的api已经集成到tensorflow 2.0新的tf.distribute.strategy中，作为 tf.distribute.strategy的一部分开源：
https://www.tensorflow.org/alpha/guide/distribute_strategy
团队还公开了相关论文：tf-replicator: distributed machine learning for researchers，全面描述了这个新框架的技术细节。
https://arxiv.org/abs/1902.00465
接下来，我们将介绍tf-replicator背后的想法和技术挑战。
构建一个分布式机器学习库
虽然tensorflow为cpu、gpu和tpu设备都提供了直接支持，但是在目标之间切换需要用户付出大量的努力。这通常涉及为特定的硬件目标专门编写代码，将研究想法限制在该平台的功能上。
一些构建在tensorflow之上的现有框架，例如estimators，已经试图解决这个问题。然而，它们通常针对生产用例，缺乏快速迭代研究思路所需的表达性和灵活性。
我们开发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)一个定义其模型逻辑的step函数(例如，梯度下降的单个step):
# deploying a model with tpureplicator.repl = tf_replicator.tpureplicator( num_workers=1, num_tpu_cores_per_worker=8)with repl.context(): model = resnet_model() base_optimizer = tf.train.adamoptimizer() optimizer = repl.wrap_optimizer(base_optimizer)# ... code to define replica input_fn and step_fn.per_replica_loss = repl.run(step_fn, input_fn)train_op = tf.reduce_mean(per_replica_loss)with tf.train.monitoredsession() as session: repl.init(session) for i in xrange(num_train_steps): session.run(train_op) repl.shutdown(session)
将计算扩展到多个设备需要设备之间进行通信。在训练机器学习模型的背景下，最常见的通信形式是累积梯度(accumulate gradients)以用于优化算法，如随机梯度下降。
因此，我们提供了一种方便的方法来封装tensorflow optimizers，以便在更新模型参数之前在设备之间累积梯度。对于更一般的通信模式，我们提供了类似于mpi的原语，如“all_reduce”和“broadcast”。这些使得实现诸如全局批标准化之类的操作变得非常简单，这是扩展biggan模型训练的关键技术。
输入数据从主机发送到各个gpu, gpu立即开始处理。当需要在gpu之间交换信息时，它们会在发送数据之前进行同步。
实现
对于多gpu计算，tf-replicator依赖于“图内复制”(“in-graph replication)模式，其中每个设备的计算在同一个tensorflow graph中复制。设备之间的通信是通过连接设备对应子图中的节点来实现的。在tf-replicator中实现这一点很具挑战性，因为在tensorflow graph中的任何位置都可能发生通信。因此，构造计算的顺序至关重要。
我们的第一个想法是在一个单独的python线程中同时构建每个设备的子图。当遇到通信原语时，线程同步，主线程插入所需的跨设备计算。之后，每个线程将继续构建其设备的计算。
然而，在我们考虑这种方法时，tensorflow的图形构建api不是线程安全的，这使得在不同线程中同时构建子图非常困难。相反，我们使用图形重写(graph rewriting)在所有设备的子图构建完成后插入通信。在构造子图时，占位符被插入到需要通信的位置。然后，我们跨设备收集所有匹配占位符，并用适当的跨设备计算替换它们。
当tf-replicator构建一个in-graph replicated计算时，它首先独立地为每个设备构建计算，并将占位符留给用户指定的跨设备计算。构建好所有设备的子图之后，tf-replicator通过用实际的跨设备计算替换占位符来连接它们。
为ai研究构建一个平台
通过在tf-replicator的设计和实现过程中与研究人员密切合作，我们最终构建一个库，让用户能够轻松地跨多个硬件加速器进行大规模计算，同时让他们拥有进行前沿ai研究所需的控制和灵活性。
例如，在与研究人员讨论之后，我们添加了mpi风格的通信原语，如all-reduce。tf-replicator和其他共享基础架构使我们能够在稳健的基础上构建越来越复杂的实验，并在整个deepmind快速传播最佳实践。
在撰写本文时，tf-replicator已经成为deepmind应用最广泛的tpu编程接口。虽然这个库本身并不局限于训练神经网络，但它最常用来训练大量数据。例如，biggan模型是在一个512核的tpuv3 pod训练的，batch size为2048。
在采用分布式actor-learner设置的增强学习智能体中，例如我们的重要性加权actor-learner架构，可扩展性是通过让许多actor通过与环境的交互生成新的体验来实现的。然后，learner对这些数据进行处理，以改进agent的策略，表示为一个神经网络。为了应对越来越多的actor，tf-replicator可以很轻松地将learner分布在多个硬件加速器上。
这些以及更多例子在我们的arxiv论文中有更详细的描述。
blog:
https://deepmind.com/blog/tf-replicator-distributed-machine-learning/
paper：
https://arxiv.org/abs/1902.00465
2019世界智能网联汽车大会将引发汽车新时代的潮流
电压跟随器的作用
抢单台积电,三星晶圆代工的底气何在?
中国最大的工程机械公司徐工，和中国最大的空调企业格力都遭遇过出售的生死危机
大话西游手游联手冰峰，国风老字号重焕江湖新情义
DeepMind又放福利：开源了一个内部的分布式机器学习库TF-Replicator
小米PK华为：智能家居这场战争谁能笑到最后
拒绝等待，华为云CDN下载加速就是要快人一步
99元的Xisem西圣ASN蓝牙耳机，真的香！
0xUniverse : 除了脚下的土地，我们还拥有广阔的星空
51单片机的汇编程序库
矿用联网机锚网气动联扣机CQ24CQ50C4530型号齐全有煤安
爆料称明年新款iPhone刘海将缩小
广和通与多家企业联合成立深圳市5G-A应用创新产业联盟
功率放大器的性能特点和使用步骤
地平线与智能汽车计算平台方案商国汽智控签署战略合作协议
LED灯珠内氯的来源及解决办法
典型案例中电源管理解决方案应用探讨
为什么一定是三次握手和四次挥手
3月新能源汽车上险数据的一些基本分析