分散式 TensorFlow：Distribution Strategy API 丨Google 開發者大會 2018

摘要： Google 開發者大會 (Google Developer Days，簡稱 GDD) 是展示 Google 最新開發者產品和平臺的全球盛會，旨在幫助你快速開發優質應用，發展和留住活躍使用者群，充分利用各種工具獲得更多收益。2018 Google 開發者大會於 9 月 20 日和 21 日於...

Google 開發者大會 (Google Developer Days，簡稱 GDD) 是展示 Google 最新開發者產品和平臺的全球盛會，旨在幫助你快速開發優質應用，發展和留住活躍使用者群，充分利用各種工具獲得更多收益。2018 Google 開發者大會於 9 月 20 日和 21 日於上海舉辦。:point_right: ofollow,noindex">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

服務並提供介面，筆記本通過介面來實現分散式訓練。

如上圖所示，啟動叢集。

程式碼例項如下：

在筆記本上啟動之後，我們可以檢視到執行日誌。可以通過關閉程序來結束訓練，也同樣可以修改模型用同樣的叢集來進行訓練。