『TensorFlow』分布式訓練_其一_邏輯梳理

1，PS-worker架構

將模型維護和訓練計算解耦合，將模型訓練分為兩個作業（job）：

• 模型相關作業，模型參數存儲、分發、匯總、更新，有由PS執行
• 訓練相關作業，包含推理計算、梯度計算（正向/反向傳播），由worker執行

該架構下，所有的woker共享PS上的參數，並按照相同的數據流圖傳播不同batch的數據，計算出不同的梯度，交由PS匯總、更新新的模型參數，大體邏輯如下：

1. pull：各個woker根據數據流圖拓撲結構從PS獲取最新的模型參數
2. feed：各個worker根據定義的規則填充各自batch的數據
3. compute：各個worker使用第一步的模型參數計算各自的batch數據，求出各自batch的梯度
4. push：各個worker將各自的梯度推送到PS
5. update：PS匯總來自n個worker的n份梯度，求出平均值后更新模型參數

分布式經典架構PS-worker會重復上面步驟，直到損失到達閾值或者輪數到達閾值。

2，數據並行模式分類

根據數據流圖構建模式分類：

• 圖內復制：單進程、‘單機多卡’的數據並行訓練，需要用戶自己實現梯度匯總和均值計算。實例，models/tutorials/image/cifar10/cifer10_multi_gpu-train.py（見下節）
• 圖間復制：多進程、跨多機的分布式訓練，使用同步優化器（SyncReplicasOptimizer）實現分布式梯度計算和模型參數更新。實例，tensorflow/tools/dist_test/python/mnist_replica.py（分布式同步訓練實踐，見下下節）

根據參數更新機制分類：

• 異步訓練：各個worker獨立訓練，計算出梯度后即刻更新參數，不需要等待其他worker完成計算
• 同步訓練：所有worker完成本輪計算后，匯總梯度，更新模型，計算能力強的worker需要阻塞等待其他worker

兩種訓練機制同時支持上面兩周數據流圖構建模式。一般來說同步機制收斂快，異步單步計算快，但易受單批數據影響，不穩定。

3，同步優化器

tensorflow進行同步（同步訓練模式專用）各個worker梯度並進行優化時，會使用特殊的優化器即同步優化器，tf.train.SyncReplicasOptimizer，其第一個參數為普通優化器，我們可以定義一個普通的優化器傳入，后續參數如下：

參數名稱	功能說明	默認值
replicas_to_aggragate	並行副本數	num_workers
total_num_replicas	實際副本數（worker數目）	num_workers

並行副本數指期望的每一步中並行的batch數據數目，實際副本數指參與的workers數目，

• 並行=實際：全民參與，一個worker領取一個batch數據
• 並行>實際：能者多勞，先完成自己batch的worker會繼續領取未訓練數據，PS會等到梯度份數到達並行數后進行模型參數計算
• 並行<實際：替補等位，存在空閑的worker，取代可能出現的異常worker，確保訓練過程高可用

運算過程

• 計算梯度過程同普通優化器，調用基類的Optimizer的compute_gradients成員方法
• 更新參數時重寫了Optimizer的apply_gradients方法，見tensorflow/python/training/sync_replicas_optimizer.py

講解同步優化器工作邏輯之前，介紹兩個概念，

梯度聚合器

每一個模型參數有一個自己隊列，收集來自不同worker的梯度值，梯度聚合器包含M個隊列對應M個模型參數，每個隊列收集來自N個worker計算出來的N個梯度值。

同步標記隊列

存儲同步標記，實際上就是N個global_step值，每個worker領取一個，用於控制同步

以全民參與模式為例

worker工作模式如下：

1. 從同步標記隊列領取一個global_step，表示全局訓練步數的同步標記
2. 將同步標記值賦予worker的本地訓練步數local_step
3. 從PS獲取最新模型參數
4. 計算出M個梯度值
5. 將M個梯度值推送到PS上的M個梯度隊列中

PS工作模式如下：

1. 從梯度聚合器上收集worker推送過來的梯度值，每個隊列收集N份（對應N個global_step下訓練值）后，計算均值，收集齊M個均值后，得到M對{模型參數，梯度值}的聚合元組
2. 更新模型參數
3. 向同步標記隊列推送N個global_step+1標記

聚合器收集梯度值並校驗local_step是否符合global_step，是則接收梯度值，計算能力強的worker提交梯度后由於沒有同步標記可以領取所以被阻塞，PS集齊N份后更新參數，發布下次N個同步標記，開始下一步訓練。

由於初始PS不會更新參數發布同步標記，所以需要初始化同步標記隊列——sync_init_op，直接向隊列注入N個0標記。

分布式模型訓練需要的主要初始化操作如下（opt指tf.train.SyncReplicasOptimizer）：

操作名稱	常用變量名	功能說明
opt.local_step_init_op	local_init_op	loacl_step初始值
pot.chief_init_op	local_init_op	gobal_step初始值
opt.ready_for_local_init_op	ready_for_local_init_op	為未初始化的Variable設置初始值
opt.get_chief_queue_runner	chief_queue_runner	同步標記隊列啟動QueueRunner實例
opt.get_init_tockens_op	sync_init_op	同步標記隊列初始化
tf.global_variables_initializer	init_op	全局Variable設置初始值

如果使用模型管理類Supervsor，可以將大部分工作交由其代勞。

以能者多勞模式對比

模型參數個數M，worker個數N，並行副本數R（R>N），此時

梯度聚合器仍然有M個參數收集隊列，每一個隊列要收集R份才進行匯總，R>N所以會存在某個worker領取多份數據的情況。

同步標記隊列存儲R個同步標記，以確保每一步中梯度聚合器可以收集到R份數據。

4，異步優化器

異步優化器沒有很多附加參量，和單機訓練幾乎一致，只是每個worker獲取參數需要從另一個進程PS中得到而已。

5，模型管理類Supervsor

本質上是對Saver（模型參數存儲恢復）、Coordinator（多線程服務生命周期管理）、SessionManager（單機以及分布式會話管理）三個類的封裝

Coordinator會監測程序的線程是否運行正常，任何異常的出現都會向Supervisor報告，此時Coordinator講程序的停止條件設置為True，Supervisor停止訓練並清理工作（關閉會話、回收內存等），其他服務檢測到True后會各自關閉服務，終止線程。

SessionManager幫助用戶創建管理單機或是分布式會話，以便簡化數據流圖的生命周期和維護邏輯，同事負責將checkpoint文件中加載出的數據恢復到數據流圖中。

流程邏輯如下：

1. 創建Supervisor實例，構造方法需要傳入checkpoint文件和summary文件存儲目錄（Supervisor的logdir參數）
2. 調用tf.train.Supervisor.managed_session，從Supervisor實例獲取會話實例
3. 使用該會話執行訓練，訓練中需要檢查停止條件，保證訓練正確性

獲取managed_session時，Supervisor會通過QueueRunner同時啟動一下三個服務：

• 檢查點服務：將數據流圖中的參數定期保存，默認10min保存一次，且會識別global_step（Supervisor的global_step參數）
• 匯總服務：默認2min一次
• 步數計數器服務：向匯總添加global_step/sec，2min一次

使用managed_session創建會話時，會自動恢復上一次的結果並繼續訓練。