數據集暴增壓力下，微信「掃一掃」識物訓練如何優雅破局？

微信「掃一掃」識物上線一段時間，由前期主要以商品圖（鞋子/箱包/美妝/服裝/家電/玩具/圖書/食品/珠寶/傢俱/其餘）做爲媒介來挖掘微信內容生態中有價值的信息，擴張到各類垂類領域的識別，包括植物/動物/汽車/果蔬/酒標/菜品/地標識別等，識別核心依託於深度學習的卷積神經網絡模型。隨着天天千萬級的增加數據和愈來愈多的模型參數量，深度學習訓練一次時間大概須要一週左右。如何可以快速訓練優化模型並上線，成爲咱們亟待解決的問題。

1、引言

現在，依託強大的GPU算力，深度學習獲得迅猛發展。在圖像處理、語音識別領域掀起了史無前例的一場革命。相較於傳統的方法，以卷積神經網絡（CNN）爲表明的深度學習方法能夠高度地重點學習數據的特性，在圖像處理領域已經取得了統治地位。

隨着掃一掃識物日調用量的持續增長，圖片數據正以天天千萬級的量級增加，在這個爭分奪秒的時代裏，得數據者得天下。同時，神經網絡的複雜性呈爆炸式增加，像15年微軟提出圖像分類的ResNet模型有7 ExaFLOPs/6千萬個參數，17年穀歌的神經網絡機器翻譯模型有100 ExaFLOPS/87億個參數。

在大部分場景下，模型能夠在一臺GPU服務器上，使用一個或者多個GPU進行訓練。但隨着數據集的增大，訓練時間也相應增加，有些時候訓練須要一週甚至更長時間。所以，如何可以快速迭代優化深度學習模型，成爲咱們算法開發者亟須解決的問題。

下文將經過從分佈式訓練方法的選擇、多機通訊技術原理進行講解，基於Horovod的框架在微信自研平臺打通分佈式訓練和實驗結果來介紹微信掃一掃識物中的深度學習模型分佈式訓練。

2、分佈式訓練

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的分佈式訓練方法。

3、多機通訊技術

相比於單機多卡，多機多卡分佈式訓練要保證多臺機器之間是能夠互相通訊的以及不一樣機器之間梯度可傳遞。

並行任務的通訊通常能夠分爲點對點通訊和集體通訊。點對點通訊這種模式只有一個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通訊。

4、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的損失規約，而後再進行梯度傳播。最後在保存模型時，只要指定一臺機器保存模型便可。

5、實驗結果

分佈式訓練除了訓練階段的通訊要儘量的快，數據的IO也是須要考慮的地方。掃一掃識物的檢索模型是基於大量的圖像數據訓練的。在進行分佈式訓練時，每一個機器都須要可以讀取這些訓練數據，圖片文件存到微信自研分佈式存儲系統上。

在訓練時，分佈式訓練的加速比和GPU數目正相關。在mnist數據集上基於resnet50測試分佈式訓練運行時間, 單機運行mnist 100個epoch須要78min，多機使用4塊GPU訓練100個epoch 23min。

在咱們實際項目的模型訓練中，基於分佈式訓練能夠將以往須要訓練5天甚至一週的時間縮短到1天之內，在一樣的時間內，算法開發者能夠探索更多的實驗，快速反饋更新，大大提升了算法研發的效率。

6、總結與展望

目前掃一掃識物在微信自研訓練平臺上可以成功進行分佈式訓練，但仍然存在如下問題：如何可以高效地存放讀取大量圖片小文件，減小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/tutor...

[3]https://developer.nvidia.com/...

[4]https://eng.uber.com/horovod/