雷锋网 AI 科技评论按:DeepMind 有一支专门的科研平台团队(the Research Platform Team),他们的职责是为 AI 学术研究构建加快计算的基本举措措施。他们不经常表态,然则此次由他们撰文介绍的 TF-Replicator 是一个极为有效的对象:它是又一个实用的软件库,可以赞助从未接触过分布式体系的研究人员们轻松地在 GPU 集群和云 TPU 集群上安排 TensorFlow 模型,也就成为了深度进修进行大年夜范围工业化应用的重要组件。TF-Replicator 的法度榜样化模型如今也已经作为 TensorFlow 的 tf.distribute.Strategy 的一部分开源在 https://www.tensorflow.org/alpha/guide/distribute_strategy

DeepMind 科研平台团队揭橥的技巧博客中介绍了 TF-Replicator 的核心思惟、开辟经历、技巧挑衅。雷锋网 AI 科技评论把重要内容编译如下。

在 AlphaFold、BigGAN、AlphaStar 等近期的人工智能科研成果中,我们很轻易看到一个反复出现的要素,那就是对轻松增长模型和计算才能范围的渴求。在越来越高的计算才能上练习越来越大年夜的模型,让神经网站展示出了惊人的表示。为了应对如许的需求,DeepMind 科研平台团队开辟了 TF-Replicator,它可以赞助科研人员轻松地为机械进修模型指定不合的硬件加快器、把运算拓展到更多设备上,以及在不合类型的加快器之间无缝切换。TF-Replicator 一开端是基于 TensorFlow 进行开辟的,如今它的 API 也已经直接集成进了 TF 2.0 中。

TensorFlow 固然对 CPU、GPU、TPU 都供给了直接的支撑,然则用户想要切换模型运行在哪个设备上倒是一件很麻烦的工作。一般来说,这个过程都须要针对目标硬件从新编写专门的代码,这也就限制了科研设法主意在具体平台上的发挥。今朝也有一些其余基于 TensorFlow 的框架在测验测验应对如许的问题,比如 Estimators,不过这些筹划根本上都是针对临盆情况的,缺乏学术科研的快速迭代所须要的表达才能和灵活性。

TF-Replicator API

构建一个分布式机械进修库

DeepMind 科研平台团队开辟 TF-Replicator 的初始目标是为 DeepMind 的研究人员们供给一种简单便捷的应用 TPU 的 API。不过 TensorFlow 原生的 TPU 计算 API 和 GPU 计算 API 有不少差别,这就为切换到 TPU 带来了不小障碍。TF-Replicator 供给的 API 要更简单、更友爱,隐蔽了 TF 本来的 TPU API 中的复杂之处。个中最关键的是,科研平台团队和 DeepMind 之外的很多机械进修组织机构都进行了慎密的合作,以确保他们开辟出的 TF-Replicator 有足够的灵活性、易用性。

应用了 TF-Replicator 的代码看起来和单台设备上运行的 TF 代码异常类似,给了用户足够的自由来定义他们本身的模型运行轮回。用户只须要做两次简单的定义,一,定义指向某个数据集的输入函数;二,定义一个指定了模型逻辑的步调函数(比如梯度降低中的某一步迭代)。代码示例如下:



# Deploying a model with TpuReplicator.

model = resnet_model()

num_workers=1, num_tpu_cores_per_worker=8

)

with repl.context():

base_optimizer = tf.train.AdamOptimizer()

optimizer = repl.wrap_optimizer(base_optimizer)

repl = tf_replicator.TpuReplicator(

# ... 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):

repl.shutdown(session)

想要把运算拓展到多个设备,须要让这些设备之间可以或许互相沟通。在练习机械进修模型的情境中,最常见的沟通情势就是为随机梯度降低 SGD 之类的优化算法累积计算梯度。所以 DeepMind 科研平台团队在 TF-Replicator 顶用一个便捷的办法供给了 TensorFlow 各类优化器的打包,如许在更新模型的参数之前,不合设备分别计算的梯度就可以先辈行累积。对于其它的更通用化的沟通,TF-Replicator 供给了类似 MPI 的主动式沟通模式,比如 `all_reduce` 和 `broadcast`。这些功能的赞助下,BigGAN 模型中应用到的全局批量归一化(global batch normalisation)也可以简单地在操作中实现,而这也是 BigGAN 增长练习范围异常重要的一步(可以拜见 https://arxiv.org/abs/1809.11096论文第三节)。

数据从主机 host 送入不合的 GPU 中并立时开端计算。当 GPU 之间须要交换信息的时刻,它们会先同步再发送数据。TF-Replicator 的实现

对于多 GPU 设备中的运算,TF-Replicator 应用的是一种“图内复制”(in-graph replication)的模式,每个设备的计算都邑复制到同一个 TensorFlow 图中。设备之间的沟通方法是连接不合设备的对应子图的节点。在 TF-Replicator 中实现这个其实很有挑衅性,因为沟通可以产生在数据流图中的任何一点上。运算的构建次序就变得尤其重要。

DeepMind 科研平台团队最开端的设法主意是在零丁的 Python 线程中为每个设备同步构建各自的子图。当碰到一个主动提议沟通的请求时,这些线程就会同步,主线程会插入履行所需的跨设备运算。在此之后,每个线程都邑持续构建它本身对应的设备的运算。然而,对这种办法进行细心思虑之后他们意识到 TF 中的图构建 API 并不是线程安然的,这意味着同步在不合的线程中构建不合的子图异常艰苦。所以他们转而应用了图重写(graph rewriting),在所有设备的子图都完成构建之后插入沟通。当这些子图还在构建的时刻,TF-Replicator 会在须要沟通的处所插入一个占位符,子图构建完成之后它就会在不合的设备间比对所有匹配的占位符,然后把它们改换成精确的跨设备运算。

当 TF-Replicator 构建一个图内复制计算时,它会起首为每个设备分别零丁构建运算,然后把用户指定了须要应用跨设备运算的处所用占位符代替。当所有设备的子图都构建完毕后,TF-Replicator 就把占位符调换成实际的跨设备运算,最终连接了不合的设备。在 DeepMind 构建一个基本平台

因为在设计和实现的全过程中都和学术界内的研究人员们保持慎密的合作,TF-Replicator 得以被构建成一个可以轻松地在多种硬件加快器之间拓展运算的库,同时还让用户保存了前沿 AI 研究所需的高度控制和灵活性。

在开辟完毕之前,TF-Replicator 就已经是 DeepMind 的研究人员们应用最多的 TPU 编程接口。TF-Replicator 的功能并不限制于练习神经收集,不过它照样最常用来在大年夜量数据长进行大年夜批量练习。比如 BigGAN 模型就是在最高达到 512 个 TPUv3 核心构成的集群上,以 2048 的批量大年夜小进行练习的。在带有分布式的履行者-进修者设置的强化进修智能体中,更大年夜范围的运算可以让更多不合的履行者和情况进行交互,也可以产生新的经验。TF-Replicator 已经为更多履行者的情况做好了预备,可以轻松地在多个硬件加快器上分布进修者(平日是神经收集)的运算。在论文中有更具体的介绍。

session.run(train_op)

DeepMind 科研平台团队开辟了很多有影响力的技巧对象,TF-Replicator 只是个中之一。DeepMind 的很多深度进修科研成果,都有科研平台团队的功绩。DeepMind 科研平台团队在介绍分享本身的成果的同时,也迎接对前沿 AI 研究感兴趣的科研平台开源软件工程师参加他们。

具体介绍请见论文 https://arxiv.org/abs/1902.00465

via DeepMind Blog,雷锋网 AI 科技评论编译

点击浏览原文,查DeepMind 开源TRFL,又一个强化进修复现、立异好副手

声明:该文观点仅代表作者本人,南宁资讯网系信息发布平台,本站仅提供信息存储空间服务。