tinroy172

DeepMind 科研平台团队发表的技术博客中介绍了 TF-Replicator 的核心思想、开发经历、技术挑战。
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)

DeepMind 科研平台团队最开始的想法是在单独的 Python 线程中为每个设备同步构建各自的子图。当遇到一个主动发起沟通的请求时，这些线程就会同步，主线程会插入执行所需的跨设备运算。在此之后，每个线程都会继续构建它自己对应的设备的运算。然而，对这种方法进行仔细思考之后他们意识到 TF 中的图构建 API 并不是线程安全的，这意味着同步在不同的线程中构建不同的子图非常困难。所以他们转而使用了图重写（graph rewriting），在所有设备的子图都完成构建之后插入沟通。当这些子图还在构建的时候，TF-Replicator 会在需要沟通的地方插入一个占位符，子图构建完成之后它就会在不同的设备间比对所有匹配的占位符，然后把它们更换成正确的跨设备运算。
在 DeepMind 构建一个基础平台
在开发完毕之前，TF-Replicator 就已经是 DeepMind 的研究人员们使用最多的 TPU 编程接口。TF-Replicator 的功能并不限制于训练神经网络，不过它还是最常用来在大量数据上进行大批量训练。比如 BigGAN 模型就是在最高达到 512 个 TPUv3 核心组成的集群上，以 2048 的批量大小进行训练的。在带有分布式的执行者-学习者设置的强化学习智能体中，更大规模的运算可以让更多不同的执行者和环境进行交互，也可以产生新的经验。TF-Replicator 已经为更多执行者的环境做好了准备，可以轻松地在多个硬件加速器上分布学习者（通常是神经网络）的运算。在论文中有更详细的介绍。
DeepMind 科研平台团队开发了许多有影响力的技术工具，TF-Replicator 只是其中之一。DeepMind 的许多深度学习科研成果，都有科研平台团队的功劳。DeepMind 科研平台团队在介绍分享自己的成果的同时，也欢迎对前沿 AI 研究感兴趣的科研平台开源软件工程师加入他们。

[ 本帖最后由 tinroy172 于 2019-3-14 11:31 编辑 ]