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