分佈式 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 技術乾貨