怎么把GPU上训练的模型转到TPU或者CPU上去？DeepMind发布新工具支招

雷锋网 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，不过这些方案基本上都是针对生产环境的，缺少学术科研的快速迭代所需要的表达能力和灵活性。

构建一个分布式机器学习库

DeepMind 科研平台团队开发 TF-Replicator 的初始目标是为 DeepMind 的研究人员们提供一种简单便捷的使用 TPU 的 API。不过 TensorFlow 原生的 TPU 计算 API 和 GPU 计算 API 有不少区别，这就为切换到 TPU 带来了不小障碍。TF-Replicator 提供的 API 要更简单、更友好，隐藏了 TF 原本的 TPU API 中的复杂之处。其中最关键的是，科研平台团队和 DeepMind 之外的许多机器学习组织机构都进行了紧密的合作，以确保他们开发出的 TF-Replicator 有足够的灵活性、易用性。

TF-Replicator API

使用了 TF-Replicator 的代码看起来和单台设备上运行的 TF 代码非常相似，给了用户足够的自由来定义他们自己的模型运行循环。用户只需要做两次简单的定义，一，定义指向某个数据集的输入函数；二，定义一个指定了模型逻辑的步骤函数（比如梯度下降中的某一步迭代）。代码示例如下：

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

想要把运算拓展到多个设备，需要让这些设备之间能够相互沟通。在训练机器学习模型的情境中，最常见的沟通形式就是为随机梯度下降 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 已经为更多执行者的环境做好了准备，可以轻松地在多个硬件加速器上分布学习者（通常是神经网络）的运算。在论文中有更详细的介绍。

DeepMind 科研平台团队开发了许多有影响力的技术工具，TF-Replicator 只是其中之一。DeepMind 的许多深度学习科研成果，都有科研平台团队的功劳。DeepMind 科研平台团队在介绍分享自己的成果的同时，也欢迎对前沿 AI 研究感兴趣的科研平台开源软件工程师加入他们。

详细介绍请见论文 https://arxiv.org/abs/1902.00465

via DeepMind Blog，雷锋网(公众号：雷锋网) AI 科技评论编译

。