微信“掃一掃”識物上線一段時間，由前期主要以商品圖（鞋子/箱包/美妝/服裝/家電/玩具/圖書/食品/珠寶/傢俱/其他）作為媒介來挖掘微信內容生態中有價值的資訊，擴張到各種垂類領域的識別，包括植物/動物/汽車/果蔬/酒標/菜品/地標識別等，識別核心依託於深度學習的卷積神經網路模型。隨著每天千萬級的增長資料和越來越多的模型引數量，深度學習訓練一次時間大概需要一週左右。如何能夠快速訓練優化模型並上線，成為我們亟待解決的問題。
一、引言
如今，依託強大的GPU算力，深度學習得到迅猛發展。在影像處理、語音識別領域掀起了前所未有的一場革命。相較於傳統的方法，以卷積神經網路（CNN）為代表的深度學習方法可以高度地重點學習資料的特性，在影像處理領域已經取得了統治地位。
隨著掃一掃識物日呼叫量的持續增加，圖片資料正以每天千萬級的量級增長，在這個爭分奪秒的時代裡，得資料者得天下。同時，神經網路的複雜性呈爆炸式增長，像15年微軟提出影像分類的ResNet模型有7 ExaFLOPs/6千萬個引數，17年穀歌的神經網路機器翻譯模型有100 ExaFLOPS/87億個引數。
在大部分場景下，模型可以在一臺GPU伺服器上，使用一個或者多個GPU進行訓練。但隨著資料集的增大，訓練時間也相應增長，有些時候訓練需要一週甚至更長時間。因此，如何能夠快速迭代優化深度學習模型，成為我們演算法開發者亟須解決的問題。
下文將通過從分散式訓練方法的選擇、多機通訊技術原理進行講解，基於Horovod的框架在微信自研平臺打通分散式訓練和實驗結果來介紹微信掃一掃識物中的深度學習模型分散式訓練。
二、分散式訓練
1. 並行方式
多機多卡相比較於單機單卡，可以將模型訓練的時間大大縮短。一般我們一臺伺服器只支援8張GPU卡，而採用分散式的多機多卡訓練方式，可以將幾十甚至幾百臺伺服器排程起來一起訓練一個模型，進一步突破模型訓練的上限。
按照分散式並行訓練方式，分散式訓練一般分為資料並行和模型並行兩種。
（1）資料並行
分散式系統中不同的GPU都有同一個模型的完全拷貝，每個GPU只獲得整個資料的不同部分，然後將所有GPU的神經網路梯度按照同步或者非同步的方式合併。
（2）模型並行
分散式系統中每個GPU使用相同的資料，所有GPU上只分布著模型的部分內容，訓練過程中交換神經網路的啟用部分。
因為模型並行各個部分存在一定的依賴，不能隨意增加GPU的數量，規模伸縮性差，在實際訓練中用的不多。而資料並行，各部分獨立，規模伸縮性好，實際訓練中更為常用，提速效果也更好。在實現性、容錯性和好的叢集利用率上，資料並行優於模型並行。
2. 系統架構
分散式訓練系統架構主要包括兩種：Parameter Server Architecture（就是常見的PS架構，引數伺服器）和Ring all-reduce Architecture。
（1）Parameter Server 架構
在PS架構中，叢集中的節點被分為兩類：parameter server和worker。其中parameter server存放模型的引數，而worker負責計算引數的梯度。
在每個迭代過程，worker從parameter sever中獲得引數，然後將計算的梯度返回給parameter server，parameter server聚合從worker傳回的梯度，然後更新引數，並將新的引數廣播給worker。
（2）Ring all-reduce 架構
在Ring all-reduce架構中，各個裝置都是worker，並且形成一個環，沒有中心節點來聚合所有worker計算的梯度。在一個迭代過程，每個worker完成自己的mini-batch訓練，計算出梯度，並將梯度傳遞給環中的下一個worker，同時它也接收從上一個worker的梯度。對於一個包含N個worker的環，各個worker需要收到其它N-1個worker的梯度後就可以更新模型引數。
採用PS計算模型的分散式，通常會遇到網路的問題，隨著worker數量的增加，其加速比會迅速的惡化。相比PS架構，Ring all-reduce架構網路通訊量不隨著worker（GPU）的增加而增加，是一個恆定值，叢集中每個節點的頻寬都被充分利用。
3. 引數更新
（1）同步更新
所有 GPU 在同一時間點與引數伺服器交換、融合梯度。在每輪訓練的時候需要彙總所有 worker訓練得到的梯度值，然後取平均值來更新引數伺服器上的模型引數。
（2）非同步更新
所有GPU 各自獨立與引數伺服器通訊，交換、融合梯度。每個 worker 在每輪訓練開始前從引數伺服器獲取模型引數，讀取訓練資料，進行訓練，訓練結束後便立即應用梯度來更新引數伺服器上的模型引數。
非同步更新通訊效率高速度快，但往往收斂不佳，因為一些速度慢的節點總會提供過時、錯誤的梯度方向。同步更新通訊效率低，通常訓練慢，但訓練收斂穩定，因為同步更新基本等同於單卡調大的batch size訓練。但是傳統的同步更新方法（各個GPU卡算好梯度，求和算平均的方式），在融合梯度時，會產生巨大的通訊資料量。
通過比對不同分散式並行方式、系統架構和引數更新，微信掃一掃識物最終選擇基於資料並行的引數同步更新的Ring all-reduce的分散式訓練方法。
三、多機通訊技術
相比於單機多卡，多機多卡分散式訓練要保證多臺機器之間是可以互相通訊的以及不同機器之間梯度可傳遞。
並行任務的通訊一般可以分為點對點通訊和集體通訊。點對點通訊這種模式只有一個sender和一個receiver，實現起來比較簡單。而涉及到分散式訓練，一般是多臺伺服器，用到集體通訊模式，包含多個sender多個receiver。集體通訊常用的通訊方式主要有下面幾個：broadcast、gather、scatter、reduce、all-reduce等。
1. MPI
在微信的自研訓練平臺中，多機的通訊是基於訊息傳遞介面(Message Passing Interface，MPI)來實現的，MPI是一種基於資訊傳遞的並行程式設計技術，定義了一組具有可移植性的程式設計介面，是一種程式設計介面標準。
在基於MPI程式設計模型中，計算是由一個或多個彼此通過呼叫庫函式進行訊息收、發通訊的程式所組成。MPI中的通訊器定義了一組能夠互相發訊息的程式。在這組程式中，每個程式會被分配一個序號，稱作秩(rank)，程式間顯性地通過指定秩來進行通訊。MPI涉及到的一些操作包括資料移動，聚集、同步等。
由於深度學習訓練引數大多在GPU上的，如果只是依靠MPI來同步引數，引數需要先從GPU搬到CPU，然後不同機器CPU之間再通訊，通訊結束之後再將引數從CPU搬到GPU，這個過程的通訊效率是很低的。所以為了提高通訊效率，在訓練的過程中使用基於nvidia開發的NCCL進行通訊。
2. NCCL
NCCL是Nvidia Collective multi-GPU Communication Library的簡稱，是Nvidia開發的能夠實現多GPU的集體通訊的庫，能夠很方便的整合至任何深度學習的訓練框架。在實現 Allreduce、Reduce、Broadcast、Allgather等方面做了很多優化，可以在PCIe、Nvlink、InfiniBand上實現較高的通訊速度。
目前NCCL1.0版本只支援單機多卡，卡之間通過PCIe、NVlink、GPU Direct P2P來通訊。NCCL 2.0會支援多機多卡，多機間通過Sockets (Ethernet)或者InfiniBand with GPU Direct RDMA通訊。
四、Horovod訓練框架
目前分散式訓練框架有許多，Horovod 是 Uber 開源的一個深度學習工具，囊括了TensorFlow, Keras, PyTorch, and Apache MXNet 這些分散式訓練框架。
並且Horovod的梯度同步和權值同步利用基於MPI和NCCL的 all-reduce演算法,而非引數伺服器架構,通訊效率更高。Horovod可以利用NVLINK、RDMA、GPUDirectRDMA、自動檢測通訊拓撲以及回退到 PCIe 和 TCP/IP 通訊這些功能。同時，將已有的訓練程式碼改成分散式訓練程式碼，改動量少，簡化分散式訓練的執行和啟動。
基於此，微信掃一掃識物選擇Horovod的分散式訓練框架，在微信自研的訓練平臺上進行訓練。
Horovod的多機通訊初始化是基於MPI的,通過MPI初始化通訊環境和程式分配。有如下幾個常用的環境引數：
-
size: 程式數量，也就是GPU數量；
-
rank：程式的唯一ID， 0-size；
-
local size: 每個worker的本地程式數量；
-
local rank: 每個worker的程式本地唯一ID。
通過這些引數來控制機器程式之間的通訊。
由於訓練採用的是資料並行這種模式，所以需要對資料進行分散式取樣。Horovod可以直接呼叫pytorch自帶的分散式取樣函式torch.utils.data.distributed.DistributedSampler。
這種方式可以適用於簡單的分散式訓練任務。但是在識物的檢索訓練過程中，我們希望dataloader可以做一些平衡取樣或者三元組取樣，上面的sampler只支援分散式取樣。
由於pytorch的DataLoader的部分初始化引數之間存在互斥關係，如果自定義了sampler,那麼這些引數batch_size、shuffle、batch_sampler、drop_last都必須使用預設值。所以我們重寫batch_sampler，將分散式的sampler作為引數傳入新構造的batch_sampler。
Horovod內部實現了廣播操作，使模型在所有工作程式中實現一致性初始化。在載入模型權重的時候，只要在rank0的機器上載入，然後使用廣播機制，就可以將引數同步到其他機器上進行權重初始化。
在訓練過程中，計算損失函式時需要涉及到allreduce操作，將所有worker的損失規約，然後再進行梯度傳播。最後在儲存模型時，只要指定一臺機器儲存模型即可。
五、實驗結果
分散式訓練除了訓練階段的通訊要儘可能的快，資料的IO也是需要考慮的地方。掃一掃識物的檢索模型是基於大量的影像資料訓練的。在進行分散式訓練時，每個機器都需要能夠讀取這些訓練資料，圖片檔案存到微信自研分散式儲存系統上。
在訓練時，分散式訓練的加速比和GPU數目正相關。在mnist資料集上基於resnet50測試分散式訓練執行時間, 單機執行mnist 100個epoch需要78min，多機使用4塊GPU訓練100個epoch 23min。
在我們實際專案的模型訓練中，基於分散式訓練可以將以往需要訓練5天甚至一週的時間縮短到1天以內，在同樣的時間內，演算法開發者可以探索更多的實驗，快速反饋更新，大大提高了演算法研發的效率。
六、總結與展望
目前掃一掃識物在微信自研訓練平臺上能夠成功進行分散式訓練，但仍然存在以下問題：如何能夠高效地存放讀取大量圖片小檔案，減少IO的耗時。道阻且長,行則將至，在後續工作中我們將針對這些問題進行探索。
參考文獻：
[1] Li M, Andersen D G, Park J W, et al. Scaling distributed machine learning with the parameter server[C]//11th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 14). 2014: 583-598.
[2]https://mpitutorial.com/tutorials/
[3]https://developer.nvidia.com/nccl
[4]https://eng.uber.com/horovod/