Press "Enter" to skip to content

分布式 TensorFlow:Distribution Strategy API

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
服务并提供接口,笔记本通过接口来实现分布式训练。

如上图所示,启动集群。
代码实例如下:

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

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注