Note: 原文發表于我的知乎專欄：算法工程師的自我修養，歡迎關注！

Methods that scale with computation are the future of AI.

—Rich Sutton, 強化學習之父

大數據時代的互聯網應用產生了大量的數據，這些數據就好比是石油，里面蘊含了大量知識等待被挖掘。深度學習就是挖掘數據中隱藏知識的利器，在許多領域都取得了非常成功的應用。然而，大量的數據使得模型的訓練變得復雜，使用多臺設備分布式訓練成了必備的選擇。

Tensorflow是目前比較流行的深度學習框架，本文著重介紹tensorflow框架是如何支持分布式訓練的。

分布式訓練策略

模型并行

所謂模型并行指的是將模型部署到很多設備上(設備可能分布在不同機器上，下同)運行，比如多個機器的GPUs。當神經網絡模型很大時，由于顯存限制，它是難以完整地跑在單個GPU上，這個時候就需要把模型分割成更小的部分，不同部分跑在不同的設備上，例如將網絡不同的層運行在不同的設備上。

由于模型分割開的各個部分之間有相互依賴關系，因此計算效率不高。所以在模型大小不算太大的情況下一般不使用模型并行。

在tensorflow的術語中，模型并行稱之為"in-graph replication"。

數據并行

數據并行在多個設備上放置相同的模型，各個設備采用不同的訓練樣本對模型訓練。每個Worker擁有模型的完整副本并且進行各自單獨的訓練。

圖1. 數據并行示例

相比較模型并行，數據并行方式能夠支持更大的訓練規模，提供更好的擴展性，因此數據并行是深度學習最常采用的分布式訓練策略。

在tensorflow的術語中，數據并行稱之為"between-graph replication"。

分布式并行模式

深度學習模型的訓練是一個迭代的過程，如圖2所示。在每一輪迭代中，前向傳播算法會根據當前參數的取值計算出在一小部分訓練數據上的預測值，然后反向傳播算法再根據損失函數計算參數的梯度并更新參數。在并行化地訓練深度學習模型時，不同設備(GPU或CPU)可以在不同訓練數據上運行這個迭代的過程，而不同并行模式的區別在于不同的參數更新方式。

圖2. 深度學習模型訓練流程圖

數據并行可以是同步的(synchronous)，也可以是異步的(asynchronous)。

異步訓練

異步訓練中，各個設備完成一個mini-batch訓練之后，不需要等待其它節點，直接去更新模型的參數。從下圖中可以看到，在每一輪迭代時，不同設備會讀取參數最新的取值，但因為不同設備讀取參數取值的時間不一樣，所以得到的值也有可能不一樣。根據當前參數的取值和隨機獲取的一小部分訓練數據，不同設備各自運行反向傳播的過程并獨立地更新參數?？梢院唵蔚卣J為異步模式就是單機模式復制了多份，每一份使用不同的訓練數據進行訓練。

圖3. 異步模式深度學習模型訓練流程圖

異步訓練總體會訓練速度會快很多，但是異步訓練的一個很嚴重的問題是梯度失效問題(stale gradients)，剛開始所有設備采用相同的參數來訓練，但是異步情況下，某個設備完成一步訓練后，可能發現模型參數已經被其它設備更新過了，此時這個設備計算出的梯度就過期了。由于梯度失效問題，異步訓練可能陷入次優解(sub-optimal training performance)。圖4中給出了一個具體的樣例來說明異步模式的問題。其中黑色曲線展示了模型的損失函數，黑色小球表示了在t0時刻參數所對應的損失函數的大小。假設兩個設備d0和d1在時間t0同時讀取了參數的取值，那么設備d0和d1計算出來的梯度都會將小黑球向左移動。假設在時間t1設備d0已經完成了反向傳播的計算并更新了參數，修改后的參數處于圖4中小灰球的位置。然而這時的設備d1并不知道參數已經被更新了，所以在時間t2時，設備d1會繼續將小球向左移動，使得小球的位置達到圖4中小白球的地方。從圖4中可以看到，當參數被調整到小白球的位置時，將無法達到最優點。

圖4. 異步模式訓練深度學習模型存在的問題示意圖

在tensorflow中異步訓練是默認的并行訓練模式。

同步訓練

所謂同步指的是所有的設備都是采用相同的模型參數來訓練，等待所有設備的mini-batch訓練完成后，收集它們的梯度后執行模型的一次參數更新。在同步模式下，所有的設備同時讀取參數的取值，并且當反向傳播算法完成之后同步更新參數的取值。單個設備不會單獨對參數進行更新，而會等待所有設備都完成反向傳播之后再統一更新參數 。

圖5. 同步模式深度學習模型訓練流程圖

同步模式相當于通過聚合多個設備上的mini-batch形成一個更大的batch來訓練模型，相對于異步模式，在同步模型下根據并行的worker數量線性增加學習速率會取得不錯的效果。如果使用tensorflow estimator接口來分布式訓練模型的話，在同步模式下需要適當減少訓練步數(相對于采用異步模式來說)，否則需要花費較長的訓練時間。Tensorflow estimator接口唯一支持的停止訓練的條件就全局訓練步數達到指定的max_steps。

Tensorflow提供了tf.train.SyncReplicasOptimizer類用于執行同步訓練。通過使用SyncReplicasOptimzer，你可以很方便的構造一個同步訓練的分布式任務。把異步訓練改造成同步訓練只需要兩步：

在原來的Optimizer上封裝SyncReplicasOptimizer，將參數更新改為同步模式；

optimizer = tf.train.SyncReplicasOptimizer(optimizer, replicas_to_aggregate=num_workers)

在MonitoredTrainingSession或者EstimatorSpec的hook中增加sync_replicas_hook：

sync_replicas_hook = optimizer.make_session_run_hook(is_chief, num_tokens=0)

小結

下圖可以一目了然地看出同步訓練與異步訓練之間的區別。

圖6. 異步訓練模式與同步訓練模式的對比

同步訓練看起來很不錯，但是實際上需要各個設備的計算能力要均衡，而且要求集群的通信也要均衡，類似于木桶效應，一個拖油瓶會嚴重拖慢訓練進度，所以同步訓練方式相對來說訓練速度會慢一些。

雖然異步模式理論上存在缺陷，但因為訓練深度學習模型時使用的隨機梯度下降本身就是梯度下降的一個近似解法，而且即使是梯度下降也無法保證達到全局最優值。在實際應用中，在相同時間內使用異步模式訓練的模型不一定比同步模式差。所以這兩種訓練模式在實踐中都有非常廣泛的應用。

分布式訓練架構

Parameter Server架構

Parameter server架構(PS架構)是深度學習最常采用的分布式訓練架構。在PS架構中，集群中的節點被分為兩類：parameter server和worker。其中parameter server存放模型的參數，而worker負責計算參數的梯度。在每個迭代過程，worker從parameter sever中獲得參數，然后將計算的梯度返回給parameter server，parameter server聚合從worker傳回的梯度，然后更新參數，并將新的參數廣播給worker。

圖7. Parameter Server架構

Ring AllReduce架構

PS架構中，當worker數量較多時，ps節點的網絡帶寬將成為系統的瓶頸。

Ring AllReduce架構中各個設備都是worker，沒有中心節點來聚合所有worker計算的梯度。Ring AllReduce算法將 device 放置在一個邏輯環路(logical ring)中。每個 device 從上行的device 接收數據，并向下行的 deivce 發送數據，因此可以充分利用每個 device 的上下行帶寬。

圖8. Ring AllReduce架構示例

使用 Ring Allreduce 算法進行某個稠密梯度的平均值的基本過程如下：

將每個設備上的梯度 tensor 切分成長度大致相等的 num_devices 個分片；

ScatterReduce 階段：通過 num_devices - 1 輪通信和相加，在每個 device 上都計算出一個 tensor 分片的和；

AllGather 階段：通過 num_devices - 1 輪通信和覆蓋，將上個階段計算出的每個 tensor 分片的和廣播到其他 device；

在每個設備上合并分片，得到梯度和，然后除以 num_devices，得到平均梯度；

以 4 個 device上的梯度求和過程為例：

ScatterReduce 階段：

圖9. Ring-AllReduce算法的ScatterReduce階段

經過 num_devices - 1 輪后，每個 device 上都有一個 tensor 分片進得到了這個分片各個 device 上的和；

AllGather 階段：

圖10. Ring-AllReduce算法的AllGather階段

經過 num_devices - 1 輪后，每個 device 上都每個 tensor 分片都得到了這個分片各個 device 上的和；

由上例可以看出，通信數據量的上限不會隨分布式規模變大而變大，一次 Ring Allreduce 中總的通信數據量是：

相比PS架構，Ring Allreduce架構是帶寬優化的，因為集群中每個節點的帶寬都被充分利用。此外，在深度學習訓練過程中，計算梯度采用BP算法，其特點是后面層的梯度先被計算，而前面層的梯度慢于前面層，Ring-allreduce架構可以充分利用這個特點，在前面層梯度計算的同時進行后面層梯度的傳遞，從而進一步減少訓練時間。Ring Allreduce的訓練速度基本上線性正比于GPUs數目(worker數)。

2017年2月百度在PaddlePaddle平臺上首次引入了ring-allreduce的架構，隨后將其提交到tensorflow的contrib package中。同年8月，Uber為tensorflow平臺開源了一個更加易用和高效的ring allreduce分布式訓練庫Horovod。最后，tensorflow官方終于也在1.11版本中支持了allreduce的分布式訓練策略CollectiveAllReduceStrategy，其跟estimator配合使用非常方便，只需要構造tf.estimator.RunConfig對象時傳入CollectiveAllReduceStrategy參數即可。

圖11. 使用CollectiveAllReduceStrategy的偽代碼

分布式tensorflow

推薦使用 TensorFlow Estimator API 來編寫分布式訓練代碼，理由如下：

開發方便，比起low level的api開發起來更加容易

可以方便地和其他的高階API結合使用，比如Dataset、FeatureColumns、Head等

模型函數model_fn的開發可以使用任意的low level函數，依然很靈活

單機和分布式代碼一致，且不需要考慮底層的硬件設施

可以比較方便地和一些分布式調度框架(e.g. xlearning)結合使用

要讓tensorflow分布式運行，首先我們需要定義一個由參與分布式計算的機器組成的集群，如下：

cluster = {'chief': ['host0:2222'],

'ps': ['host1:2222', 'host2:2222'],

'worker': ['host3:2222', 'host4:2222', 'host5:2222']}

集群中一般有多個worker，需要指定其中一個worker為主節點(cheif)，chief節點會執行一些額外的工作，比如模型導出之類的。在PS分布式架構環境中，還需要定義ps節點。

要運行分布式Estimator模型，只需要設置好TF_CONFIG環境變量即可，可參考如下代碼：

# Example of non-chief node:

os.environ['TF_CONFIG'] = json.dumps(

{'cluster': cluster,

'task': {'type': 'worker', 'index': 1}})

# Example of chief node:

os.environ['TF_CONFIG'] = json.dumps(

{'cluster': cluster,

'task': {'type': 'chief', 'index': 0}})

# Example of evaluator node (evaluator is not part of training cluster)

os.environ['TF_CONFIG'] = json.dumps(

{'cluster': cluster,

'task': {'type': 'evaluator', 'index': 0}})

定義好上述環境變量后，調用tf.estimator.train_and_evaluate即可開始分布式訓練和評估，其他部分的代碼跟開發單機的程序是一樣的，可以參考下面的資料：

參考資料

總結

以上是生活随笔為你收集整理的tensowflow 训练 远程提交_一文说清楚Tensorflow分布式训练必备知识的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。