分布式 TensorFlow：Distribution Strategy API 丨Google 开发者大会 2018

Google 开发者大会 (Google Developer Days，简称 GDD) 是展示 Google 最新开发者产品和平台的全球盛会，旨在帮助你快速开发优质应用，发展和留住活跃用户群，充分利用各种 工具 获得更多收益。2018 Google 开发者大会于 9 月 20 日和 21 日于上海举办。:point_right: Google 开发者大会 2018 掘金专题

2018 年 9 月 21 日 周玥枫（Google Brain 软件工程师）带来一场《分布式 TensorFlow：Distribution Strategy API》的演讲，本文将对演讲做一个回顾。

为什么要分布式训练？

如上图所示，在单个 GPU 训练 Accuracy 需要花费 87 个小时。但是在实际的生产环境中，我们花费如此多的时间在单个模型的训练上。如果我们把时间从几天压缩到几个小时的话，在同样的时间里，可以训练更多的模型，尝试更多的想法，生产力将大大提高，所以使用分布式训练是实际生产环境所需要的。

Distribution Strategy API

目标：

• 简单易用 - 极少的代码修改
• 优秀的开箱性能
• 多功能 - 支持不同的分布式架构和 API

用Keras 训练 ResNet50 (非分布式)

train_dataset = tf.data.Dataset(...)
eval_dataset = tf.data.Dataset(...)
model = tf.keras.applications.ReyNet50()
optimizer = tf.train.GradientDescentOptimizer(learning_rate = 0.1)

model.compile(loss = "categorcial_crossentropy",optimizer = optimizer)

model.fit(train_dataset, epochs = 10)
model.evaluate(eval_dataset)
复制代码

用Keras 在多个 GPU 上训练 ResNet50 （在 TensorFlow r1.11 中可用）

train_dataset = tf.data.Dataset(...)
eval_dataset = tf.data.Dataset(...)
model = tf.keras.applications.ReyNet50()
optimizer = tf.train.GradientDescentOptimizer(learning_rate = 0.1)

strategy = tf.contrib.distribute.MirroredStrategy()
model.compile(loss = "categorcial_crossentropy",optimizer = optimizer, 
    distribute = strategy)

model.fit(train_dataset, epochs = 10)
model.evaluate(eval_dataset)
复制代码

通过对比我们可以发现，仅仅只是对两行代码的修改就能实现多 GPU 的训练。即使用 MirroredStrategy API 无需修改 input pipeline ，无需修改模型、训练循环就可以实现多 GPU 训练，并且无缝支持 Checkpoints，metrics, summaries.

MirroredStrategy

仅仅通过两行代码的修改，就能实现多 GPU 训练， MirroredStrategy 做了些什么？

• 数据并行
• Mirrored Variables(镜像变量)
• ALL- reduce
• 同步训练

数据并行

Mirrored Variables(镜像变量)

ALL - reduce

同步训练

• 每个设备都有模型 （模型为两层）
• 每个设备拿到数据的一部分，用本地镜像变量进行计算
• 用 ALL - reduce 交换梯度，对梯度就行求和，并用求和结果来更新本地的镜像变量
• 只有当镜像更新之后才会进行下一步，从而保证了一致性

多 GPU 性能

在 Cloud TPU 上用 keras 训练 ResNet50 (Coming up in TF r1.12)

同样只需要修改上图中方框标记的两行代码，就可以实现云端 TPU 的拓展。

TPUStrategy

拓展到多机

部署到集群

推荐使用 kubernetes 部署集群

Collective ALL-Reduce Strategy

Collective Ops具有以下几个特点：

• 自动选择最佳算法，融合梯度
• 支持插件，可以添加支持新的架构和算法
• 容错（Future）

同时，Collective 支持多种 All - Reduce 算法。

Ring All - Reduce

Hierarchical All - Reduce

其他多机架构

实现示例如下：

同样只需要修改少量的代码，就可实现我们想要的效果。

可以设置 PS 数量

新的分布式训练方法

服务并提供接口，笔记本通过接口来实现分布式训练。

如上图所示，启动集群。

代码实例如下：

在笔记本上启动之后，我们可以查看到运行日志。可以通过关闭进程来结束训练，也同样可以修改模型用同样的集群来进行训练。