作者簡介

CW，廣東深圳人，畢業(yè)于中山大學(xué)（SYSU）數(shù)據(jù)科學(xué)與計(jì)算機(jī)學(xué)院，畢業(yè)后就業(yè)于騰訊計(jì)算機(jī)系統(tǒng)有限公司技術(shù)工程與事業(yè)群（TEG）從事Devops工作，期間在AI LAB實(shí)習(xí)過，實(shí)操過道路交通元素與醫(yī)療病例圖像分割、視頻實(shí)時人臉檢測與表情識別、OCR等項(xiàng)目。

目前也有在一些自媒體平臺上參與外包項(xiàng)目的研發(fā)工作，項(xiàng)目專注于CV領(lǐng)域（傳統(tǒng)圖像處理與深度學(xué)習(xí)方向均有）。

Foreword

使用多GPU卡訓(xùn)練的情況下Batch Normalization（BN）可能會帶來很多問題，目前在很多深度學(xué)習(xí)框架如?Caffe、MXNet、TensorFlow 和?PyTorch 等，所實(shí)現(xiàn)的?BN 都是非同步的(unsynchronized)，即歸一化操作是基于每個?GPU上的數(shù)據(jù)獨(dú)立進(jìn)行的。

本文會為大家解析?BN 的多卡同步版本，這里簡稱?SyncBN，首先解釋為何需要進(jìn)行同步，接著為大家揭曉需要同步哪些信息，最后結(jié)合基于 Pytorch 實(shí)現(xiàn)的代碼解析實(shí)現(xiàn)過程中的關(guān)鍵部分。

Outline

i ?Why Synchronize BN：為何在多卡訓(xùn)練的情況下需要對BN進(jìn)行同步？

ii What is Synchronized BN：什么是同步的BN，具體同步哪些東西？

iii ?How to implement：如何實(shí)現(xiàn)多卡同步的BN？

???1. ?2次同步 vs 1次同步；

? ?2. ?介紹torch.nn.DataParallel的前向反饋；

? ?3. ?重載torch.nn.DataParallel.replicate方法；

? ?4. ?SyncBN 的同步注冊機(jī)制；

? ?5. ?SyncBN 的前向反饋

1.Why Synchronize BN：

為何在多卡訓(xùn)練的情況下需要對BN進(jìn)行同步？

對于視覺分類和目標(biāo)檢測等這類任務(wù)，batch size 通常較大，因此在訓(xùn)練時使用 BN 沒太大必要進(jìn)行多卡同步，同步反而會由于GPU之間的通信而導(dǎo)致訓(xùn)練速度減慢；

然而，對于語義分割等這類稠密估計(jì)問題而言，分辨率高通常會得到更好的效果，這就需要消耗更多的GPU內(nèi)存，因此其 batch size 通常較小，那么每張卡計(jì)算得到的統(tǒng)計(jì)量可能與整體數(shù)據(jù)樣本具有較大差異，這時候使用 BN 就有一定必要性進(jìn)行多卡同步了。

這里再提一點(diǎn)，如果使用pytorch的torch.nn.DataParallel，由于數(shù)據(jù)被可使用的GPU卡分割（通常是均分），因此每張卡上 BN 層的batch size（批次大小）實(shí)際為

，下文也以torch.nn.DataParallel為背景進(jìn)行說明。

2.What is Synchronized BN：

什么是同步的BN，具體同步哪些東西？

由開篇至今，CW 一直提到“同步”這兩個字眼，那么到底是什么是同步的BN，具體同步的是什么東西呢？

同步是發(fā)生在各個GPU之間的，需要同步的東西必然是它們互不相同的東西，那到底是什么呢？或許你會說是它們拿到的數(shù)據(jù)，嗯，沒錯，但肯定不能把數(shù)據(jù)同步成一樣的了，不然這就和單卡訓(xùn)練沒差別了，浪費(fèi)了多張卡的資源...

現(xiàn)在，聰明的你肯定已經(jīng)知道了，需要同步的是每張卡上計(jì)算的統(tǒng)計(jì)量，即?BN 層用到的

（均值）和

（方差），這樣子每張卡對其拿到的數(shù)據(jù)進(jìn)行歸一化后的效果才能與單卡情況下對一個 batch 的數(shù)據(jù)歸一化后的效果相當(dāng)。

因此，同步的?BN，指的就是每張卡上對應(yīng)的?BN 層，分別計(jì)算出相應(yīng)的統(tǒng)計(jì)量

和

，接著基于每張卡的計(jì)算結(jié)果計(jì)算出統(tǒng)一的

?和

，然后相互進(jìn)行同步，最后它們使用的都是同樣的

和

3.How to implement：

如何實(shí)現(xiàn)多卡同步的BN？

1.? 2次同步?vs 1次同步

我們已經(jīng)知道，在前向反饋過程中各卡需要同步均值和方差，從而計(jì)算出全局的統(tǒng)計(jì)量，或許大家第一時間想到的方式是先同步各卡的均值，計(jì)算出全局的均值，然后同步給各卡，接著各卡同步計(jì)算方差...這種方式當(dāng)然沒錯，但是需要進(jìn)行2次同步，而同步是需要消耗資源并且影響模型訓(xùn)練速度的，那么，是否能夠僅用1次同步呢？

全局的均值很容易通過同步計(jì)算得出，因此我們來看看方差的計(jì)算：

方差的計(jì)算，其中m為各GPU卡拿到的數(shù)據(jù)批次大小

由上可知，每張卡計(jì)算出

和

，然后進(jìn)行同步求和，即可計(jì)算出全局的方差。同時，全局的均值可通過各卡的

同步求和得到，這樣，僅通過1次同步，便可完成全局均值及方差的計(jì)算。

?2.??介紹nn.DataParallel的前向反饋

熟悉 pytorch 的朋友們應(yīng)該知道，在進(jìn)行GPU多卡訓(xùn)練的場景中，通常會使用nn.DataParallel來包裝網(wǎng)絡(luò)模型，它會將模型在每張卡上面都復(fù)制一份，從而實(shí)現(xiàn)并行訓(xùn)練。這里我自定義了一個類繼承nn.DataParallel，用它來包裝SyncBN，并且重載了nn.DataParallel的部分操作，因此需要先簡單說明下nn.DataParallel的前向反饋涉及到的一些操作。

nn.DataParallel的使用，其中DEV_IDS是可用的各GPU卡的id，模型會被復(fù)制到這些id對應(yīng)的各個GPU上，DEV是主卡，最終反向傳播的梯度會被匯聚到主卡統(tǒng)一計(jì)算。

先來看看nn.DataParallel的前向反饋方法的源碼：

其中，主要涉及調(diào)用了以下4個方法：

(1) scatter：將輸入數(shù)據(jù)及參數(shù)均分到每張卡上；

(2) replicate：將模型在每張卡上復(fù)制一份（注意，卡上必須有scatter分割的數(shù)據(jù)存在！）；

(3) parallel_apply：每張卡并行計(jì)算結(jié)果，這里會調(diào)用被包裝的具體模型的前向反饋操作（在我們這里就是會調(diào)用 SyncBN 的前向反饋方法）；

(4)? gather：將每張卡的計(jì)算結(jié)果統(tǒng)一匯聚到主卡。

注意，我們的關(guān)鍵在于重載replicate方法，原生的該方法只是將模型在每張卡上復(fù)制一份，并且沒有建立起聯(lián)系，而我們的?SyncBN?是需要進(jìn)行同步的，因此需要重載該方法，讓各張卡上的SyncBN?通過某種數(shù)據(jù)結(jié)構(gòu)和同步機(jī)制建立起聯(lián)系。

3.??重載nn.DataParallel.replicate方法

在這里，可以設(shè)計(jì)一個繼承nn.DataParallel的子類DataParallelWithCallBack，重載了replicate方法，子類的該方法先是調(diào)用父類的replicate方法，然后調(diào)用一個自定義的回調(diào)函數(shù)（這也是之所以命名為DataParallelWithCallBack的原因），該回調(diào)函數(shù)用于將各卡對應(yīng)的?SyncBN?層關(guān)聯(lián)起來，使得它們可以通過某種數(shù)據(jù)結(jié)構(gòu)進(jìn)行通信。

自定義的回調(diào)函數(shù)，將各卡對應(yīng)的Syn-BN層進(jìn)行關(guān)聯(lián)，其中DataParallelContext是一個自定義類，其中沒有定義實(shí)質(zhì)性的東西，作為一個上下文數(shù)據(jù)結(jié)構(gòu)，實(shí)例化這個類的對象主要用于將各個卡上對應(yīng)的Syn-BN層進(jìn)行關(guān)聯(lián)；_sync_replicas是在Syn-BN中定義的方法，在該方法中其余子卡上的Syn-BN層會向主卡進(jìn)行注冊，使得主卡能夠通過某種數(shù)據(jù)結(jié)構(gòu)和各卡進(jìn)行通信。

4.? Syn-BN的同步注冊機(jī)制

由上可知，我們需要在 SyncBN 中實(shí)現(xiàn)一個用于同步的注冊方法，SyncBN 中還需要設(shè)置一個用于管理同步的對象（下圖中的?_sync_master），這個對象有一個注冊方法，可將子卡注冊到其主卡。

在 SyncBN 的方法中，若是主卡，則將上下文管理器的 sync_master 屬性設(shè)置為這個管理同步的對象（_sync_master）；否則，則調(diào)用上下文對象的同步管理對象的注冊方法，將該卡向其主卡進(jìn)行注冊。

5.? Syn-BN的前向反饋

如果你認(rèn)真看完了以上部分，相信這部分你也知道大致是怎樣一個流程了。

首先，每張卡上的?SyncBN 各自計(jì)算出?mini-batch 的和

以及平方和

，然后主卡上的 SyncBN 收集來自各個子卡的計(jì)算結(jié)果，從而計(jì)算出全局的均值和方差，接著發(fā)放回各個子卡，最后各子卡的 SyncBN 收到來自主卡返回的計(jì)算結(jié)果各自進(jìn)行歸一化（和縮放平移）操作。當(dāng)然，主卡上的 SyncBN 計(jì)算出全局統(tǒng)計(jì)量后就可以進(jìn)行它的歸一化（和縮放平移）操作了。

4.最后

在同步過程中，還涉及線程和條件對象的使用，這里就不展開敘述了，感興趣的朋友可以到SyncBN源碼鏈接：https://github.com/chrisway613/Synchronized-BatchNormalization。另外，在信息同步這部分，還可以設(shè)計(jì)其它方式進(jìn)行優(yōu)化，如果你有更好的意見，還請積極反饋，CW熱烈歡迎！