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

Google 开发者大会 (Google Developer Days，简称 GDD) 是展现 Google 最新开发者产品和平台的全球盛会，旨在帮助你快速开发优质应用，发展和留住活跃用户群，充分利用各类工具得到更多收益。2018 Google 开发者大会于 9 月 20 日和 21 日于上海举办。👉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

其余多机架构

实现示例以下：

一样只须要修改少许的代码，就可实现咱们想要的效果。

经过 Kubernetes 能够设置 PS 数量

新的分布式训练方法

咱们能够在笔记本上构建咱们的模型，集群跑 TensorFlow 服务并提供接口，笔记本经过接口来实现分布式训练。

如上图所示，启动集群。

代码实例以下：

在笔记本上启动以后，咱们能够查看到运行日志。能够经过关闭进程来结束训练，也一样能够修改模型用一样的集群来进行训练。

正在实现

参考资源

以上就是本次演讲所有内容，但愿对你们有所帮助。 阅读更多 Google 开发者大会 2018 技术干货