練到推理的完整指南)
Batch Normalization這個技術(shù)我從2015年它被提出時就開始在項目里用了到現(xiàn)在快十年說它是過去十年深度學(xué)習(xí)訓(xùn)練里最“潤物細(xì)無聲”的改進(jìn)之一一點都不夸張。今天這篇不是教科書式復(fù)述論文而是結(jié)合我這些年真正在圖像分類、目標(biāo)檢測、語義分割等多個項目里反復(fù)調(diào)參、踩坑、驗證后的理解把Batch Normalization從原理到工程實現(xiàn)再到各種坑完整地拆一遍。適合剛接觸深度學(xué)習(xí)、只會調(diào)包不太懂原理的入門者也適合已經(jīng)用過一段時間但總在某些邊界場景下翻車的中級開發(fā)者甚至對老手來說后面關(guān)于BN在推理階段和遷移學(xué)習(xí)里的坑也可能幫你解釋一些之前“明明訓(xùn)練很好但測試就是崩了”的詭異現(xiàn)象。1. 為什么需要Batch Normalization從訓(xùn)練困境說起1.1 沒有BN時深層網(wǎng)絡(luò)為什么難訓(xùn)2015年之前訓(xùn)練一個深度卷積網(wǎng)絡(luò)是件非?!翱茨槨钡氖?。網(wǎng)絡(luò)層數(shù)一旦加深訓(xùn)練過程就會變得極其脆弱你精心調(diào)好的學(xué)習(xí)率在某一層可能剛剛好到另一層就變成了發(fā)散或死寂。根源在于數(shù)據(jù)在網(wǎng)絡(luò)中一層一層流動時每一層的輸出分布都在不斷變化——前一層的參數(shù)更新了后一層收到的輸入分布就整體平移或縮放后一層就得重新去適應(yīng)這個新分布。這么反復(fù)折騰的后果就是訓(xùn)練收斂慢得讓人發(fā)瘋對初始化極其敏感學(xué)習(xí)率稍微調(diào)大一點就梯度爆炸或者產(chǎn)生NaN。我記得早期做某個圖像分類Demo時換一個初始化種子同樣的代碼、同樣的超參數(shù)訓(xùn)練曲線能差出一大截。這種不確定性讓實驗對比變得非常困難——你沒法判斷模型效果的差異是來自算法改進(jìn)還是單純因為這次初始化運(yùn)氣好。一個直覺的解決思路是手工做數(shù)據(jù)標(biāo)準(zhǔn)化把輸入圖像的像素除以255或者做Z-score標(biāo)準(zhǔn)化。但這樣的標(biāo)準(zhǔn)化只在網(wǎng)絡(luò)入口做了一次網(wǎng)絡(luò)中間層照樣是“野路子”沒有任何約束每層的激活值分布依舊會隨著訓(xùn)練漂移。問題不在入口而在深層網(wǎng)絡(luò)內(nèi)部。1.2 “內(nèi)部協(xié)變量偏移”到底怎么理解論文里對BN動機(jī)的表述是“Internal Covariate Shift”中文翻譯成“內(nèi)部協(xié)變量偏移”。這個術(shù)語看著唬人其實理解起來并不難。協(xié)變量偏移原本是統(tǒng)計學(xué)里的概念訓(xùn)練數(shù)據(jù)的輸入分布跟測試數(shù)據(jù)的輸入分布不一致。而“內(nèi)部”協(xié)變量偏移指的就是神經(jīng)網(wǎng)絡(luò)內(nèi)部每一層的輸入分布在訓(xùn)練過程中隨著前面層數(shù)參數(shù)的更新而不斷變化。我習(xí)慣用一個類比來理解假設(shè)你要訓(xùn)練一批工人完成流水線作業(yè)每個工位只處理上一工位傳來的零件。上一工位的操作不穩(wěn)定零件時大時小、時歪時正下一工位就不得不花費(fèi)大量精力去適應(yīng)這些“變形零件”真正用在打磨自己工序上的精力就變少了。BN做的事情就是在每個工位入口加了一臺“標(biāo)準(zhǔn)化機(jī)”不管零件進(jìn)來什么樣先統(tǒng)一成標(biāo)準(zhǔn)規(guī)格再送進(jìn)去加工。這樣每個工位只需要在自己的職責(zé)范圍內(nèi)穩(wěn)定輸出整體流水線的效率自然就上去了。更嚴(yán)謹(jǐn)?shù)卣f我們期望網(wǎng)絡(luò)在訓(xùn)練過程中每一層都只負(fù)責(zé)“學(xué)”它應(yīng)該學(xué)的東西而不是被迫去追蹤前面層參數(shù)的“余波”。BN通過在每一層入口強(qiáng)制施加一個穩(wěn)定的統(tǒng)計分布把前面層參數(shù)變化帶來的“余波”攔截在層與層之間各個層之間就實現(xiàn)了某種程度的解耦。1.3 BN真正帶來的是什么很多人以為BN只是加快了收斂速度實際上它的貢獻(xiàn)是全方位的。第一收斂速度的提升非??捎^。我自己在多個CNN項目里對比過加了BN之后通常能用原來2到5倍的學(xué)習(xí)率起步訓(xùn)練到同樣精度所需的epoch數(shù)常常能減少三分之一甚至一半。原因很簡單歸一化之后的特征尺度被約束在固定范圍梯度傳播途中不容易因為尺度問題產(chǎn)生爆炸或消失。第二對初始化不再那么“挑食”。早期的網(wǎng)絡(luò)對權(quán)重初始化非常敏感初始化小了梯度消失初始化大了梯度爆炸。BN把每一層輸出拉回到均值為0、方差為1的尺度上相當(dāng)于給初始化兜了底。過去需要小心翼翼按層數(shù)調(diào)整初始化方差現(xiàn)在用常見的初始化方式配合BN基本都能正常訓(xùn)練起來。第三它天然帶有輕微的正則化效果。因為每個batch的統(tǒng)計量都有隨機(jī)波動模型被迫學(xué)會對擾動具有一定的魯棒性這跟Dropout的機(jī)制有相似之處。實踐中能觀察到加了BN后對Dropout的依賴會下降有些結(jié)構(gòu)甚至直接去掉Dropout也能保持差不多的泛化性能。第四它讓大學(xué)習(xí)率成為可能。深層的梯度尺度與層數(shù)息息相關(guān)沒有BN時學(xué)習(xí)率稍大就會出現(xiàn)loss直接飆到幾千萬甚至NaN的慘劇。有了BN梯度的尺度被約束在合理范圍很多人訓(xùn)練新模型時習(xí)慣性地先把學(xué)習(xí)率設(shè)成0.1這在很多場景下是可行的這在BN出現(xiàn)之前簡直是自殺行為。2. 算法原理與公式逐項拆解2.1 BN的計算流程四個步驟BN在訓(xùn)練時做的事情對每一個特征維度注意是每個維度單獨處理分四步走第一步計算當(dāng)前mini-batch內(nèi)該維度的均值$$mu_B \frac{1}{m} \sum_{i1}^{m} x_i$$第二步計算當(dāng)前mini-batch內(nèi)該維度的方差。這里要用無偏估計嗎訓(xùn)練時BN用的其實是有偏方差也就是除以m而不是m-1。這一點跟統(tǒng)計學(xué)里的樣本方差習(xí)慣不同后面細(xì)說。第三步用均值和方差把每個樣本的該維度歸一化$$\hat{x}_i \frac{x_i - mu_B}{\sqrt{\sigma_B^2 \epsilon}}$$第四步做仿射變換把歸一化后的值從標(biāo)準(zhǔn)正態(tài)分布“放回”到一個能夠表達(dá)任意分布的形狀$$y_i \gamma \hat{x}_i \beta$$其中$\gamma$和$\beta$是可學(xué)習(xí)的參數(shù)初始值通常分別設(shè)為1和0。為什么在歸一化之后還要再學(xué)一組縮放和偏移因為如果模型本來就不需要輸入是標(biāo)準(zhǔn)正態(tài)分布強(qiáng)制標(biāo)準(zhǔn)化反而會損失表達(dá)能力。比如某層學(xué)出來的特征本來就應(yīng)該分布在0到1之間且方差很小那么通過$\gamma$和$\beta$就能幾乎還原成原來的分布。換句話說理論上網(wǎng)絡(luò)可以學(xué)會“完全不使用BN”把$\gamma$設(shè)成$\sqrt{\sigma^2\epsilon}$、$\beta$設(shè)成$\mu$就能反向恢復(fù)原始輸入。這個設(shè)計保證了BN只會增加網(wǎng)絡(luò)的表達(dá)能力不會強(qiáng)制性地削弱信息。2.2 ε、γ、β每個參數(shù)到底在干什么很多初學(xué)者對著公式會覺得參數(shù)很多、不知道每個在干嘛。這里逐個說清楚。ε是一個極小的常數(shù)常見取值是1e-5。它存在的原因非常樸素防止方差為0導(dǎo)致除以0。當(dāng)一個batch內(nèi)某個維度的取值完全相同比如全連接層后的某個神經(jīng)元輸出恒定時方差恰好為0此時不加上ε就沒有數(shù)學(xué)定義。它是數(shù)值穩(wěn)定項不是超參數(shù)幾乎沒有人去調(diào)它保持默認(rèn)足夠。γ和β則是真正的可學(xué)習(xí)參數(shù)它們跟普通權(quán)重一樣通過反向傳播更新。它們存在的意義我在2.1里講過就是恢復(fù)歸一化可能“誤傷”的表達(dá)能力。實際中我觀察過不少模型的γ在訓(xùn)練后并不會乖乖停留在1附近很多通道的γ會被學(xué)到遠(yuǎn)大于1或小于1的值這說明網(wǎng)絡(luò)確實在利用這個自由度去重新縮放各通道的重要性。這里還要澄清一個常見的誤解不少人以為BN一定會把所有激活值限制在0到某個范圍因此把它當(dāng)成一種“防止梯度消失”的固定手段。實際上經(jīng)過γ和β的仿射變換后激活值完全可能被推到很大的區(qū)間BN做得只是控制分布的形狀讓它在“標(biāo)準(zhǔn)正態(tài)”和“任意分布”之間找到一個可學(xué)習(xí)的平衡點。2.3 反向傳播時BN層的梯度怎么流動因為BN引入了求均值、求方差的不可微環(huán)節(jié)反向傳播時需要把梯度回傳到每個輸入樣本上。推導(dǎo)過程這里不展開只說結(jié)論性的直覺。歸一化操作$\hat{x} (x - \mu)/\sigma$在反向傳播時對每個$x_i$的梯度不僅來自于自身歸一化值的梯度還來自于對均值$\mu$的梯度貢獻(xiàn)和對方差$\sigma^2$的梯度貢獻(xiàn)。這就產(chǎn)生了一個“間接路徑”當(dāng)某一個$x_i$變化時不僅它自己的$\hat{x}_i$變了整個batch的均值和方差也會變進(jìn)而影響batch里所有其他$x_j$的歸一化結(jié)果。所以BN在反向傳播時會形成類似“同batch相互關(guān)聯(lián)”的梯度結(jié)構(gòu)。這其實也是BN“正則化效果”的一個來源每個樣本的梯度計算都隱式地依賴batch中其他樣本等于引入了一種類對比的機(jī)制。如果某個樣本是離群點它一方面通過自身影響梯度另一方面也把均值和方差拖偏從而影響同batch內(nèi)其他樣本的梯度。模型在這種“互相牽制”中發(fā)現(xiàn)穩(wěn)定特征的能力變強(qiáng)了泛化自然更好。另外γ和β的梯度計算很輕量γ的梯度就是在歸一化值上的梯度的加權(quán)和β的梯度就是歸一化值梯度的和。它們都是按batch內(nèi)累積再平均跟普通參數(shù)的梯度更新方式一致。3. 訓(xùn)練與推理兩個完全不同的狀態(tài)3.1 訓(xùn)練時到底用的是哪種統(tǒng)計量訓(xùn)練時BN使用的是當(dāng)前mini-batch的均值與方差這是一個動態(tài)變化的統(tǒng)計量。每來一個新batch統(tǒng)計量就重新計算一次所以模型看到的分布其實一直在輕微波動。這種波動把“噪聲”注入到訓(xùn)練過程中間接提升了泛化能力。但是訓(xùn)練完成后部署模型做推理時如果仍然實時計算batch的均值和方差就會遇到兩個明顯問題。第一推理時樣本常常是一個個進(jìn)來的在線服務(wù)場景batch_size1時單個樣本的均值和方差沒有任何統(tǒng)計意義除以一個幾乎為0的方差會直接導(dǎo)致輸出爆掉。第二即使推理時有足夠大的batch也不應(yīng)該讓預(yù)測結(jié)果依賴于同batch里其他樣本——同一樣本在不同batch里預(yù)測值應(yīng)該是一致的。生產(chǎn)環(huán)境的模型預(yù)測必須是確定性的。因此BN在訓(xùn)練過程中會額外維護(hù)一組全局統(tǒng)計量滑動均值running mean和滑動方差running var。每處理一個batch就按如下方式更新一次running_mean momentum * running_mean (1 - momentum) * batch_mean running_var momentum * running_var (1 - momentum) * batch_varmomentum通常取0.9或0.99它控制全局統(tǒng)計量跟隨當(dāng)前batch統(tǒng)計量的“更新速度”。momentum越接近1統(tǒng)計量更新越慢越接近0更新越快但也會越容易受最近batch的擾動。3.2 推理公式與訓(xùn)練公式的區(qū)別推理時BN不再使用當(dāng)前batch的均值和方差而是使用訓(xùn)練階段累積好的全局統(tǒng)計量同時γ和β用訓(xùn)練好的固定值。計算過程變成y gamma * (x - running_mean) / sqrt(running_var epsilon) beta這里要強(qiáng)調(diào)一個容易忽略的細(xì)節(jié)訓(xùn)練時的歸一化用的是當(dāng)前batch的方差而且是有偏方差除以batch大小m。而推理時的running_var在絕大多數(shù)框架的實現(xiàn)里維護(hù)的是一個無偏估計實際上這里各框架的處理并不一致。以PyTorch框架為例它的BatchNorm2d在訓(xùn)練時用的是有偏方差做歸一化。更新running_var時乘以的是m/(m-1)校正因子所以running_var趨向的是無偏方差。這就造成了一個很小的偏差訓(xùn)練時歸一化用的方差偏小有偏方差通常比無偏方差小m/(m-1)倍推理時用的running_var偏大一點。這個偏差在batch足夠大時不明顯但batch_size很小比如2或4時會影響訓(xùn)練和推理之間的行為一致性。這也是為什么小batch_size下BN訓(xùn)練出來的模型測試性能往往比大batch_size更差的原因之一。3.3 eval模式下最常見的“翻車現(xiàn)場”我在實際項目里見到過太多人在這里翻車說幾個典型。最常見的坑訓(xùn)練流程里忘了切換模型的訓(xùn)練/推理狀態(tài)。尤其在一些自定義訓(xùn)練循環(huán)里有人把model.eval()漏掉了模型在推理時仍然使用當(dāng)前batch的統(tǒng)計量導(dǎo)致輸出跟期望的不一致甚至出現(xiàn)“同樣的輸入換個batch就出不同結(jié)果”的詭異behavior。這個bug在圖像分類這種“測試單張圖”的場景下很難自查因為單張圖時batch1用當(dāng)前batch統(tǒng)計量歸一化輸出往往還是能給出一個“看起來合理”的結(jié)果于是問題被掩蓋。第二個坑在PyTorch框架里BatchNorm的track_running_stats參數(shù)。默認(rèn)是True表示會自動維護(hù)running_mean和running_var。但如果有人把它設(shè)成了False那么即使在eval模式下BN也會使用當(dāng)前batch的統(tǒng)計量模型在單樣本推理時行為會完全不可控。這個參數(shù)平時很少碰但一旦碰了帶來的困惑相當(dāng)大。第三個坑模型partial加載參數(shù)時只加載了state_dict里的weight和bias忽略了running_mean和running_var。預(yù)訓(xùn)練模型做fine-tune還好world模型做推理卻忘了加載統(tǒng)計量就會出現(xiàn)“訓(xùn)練時loss正常加載模型做測試時輸出全亂套”的現(xiàn)象。這通常是state_dict里缺失了num_batches_tracked等鍵但框架不報錯只給出Missing key警告被很多人忽略掉了。第四個坑相對高級多卡訓(xùn)練時BN統(tǒng)計量的同步問題。DataParallel和DistributedDataParallel的默認(rèn)行為不同DataParallel下每個卡的BN統(tǒng)計量是各自計算的沒有跨卡同步而DDP默認(rèn)會把BN的統(tǒng)計量進(jìn)行全局同步syncBN。理論上syncBN能降低batch_size太小帶來的統(tǒng)計量噪聲但它也引入了跨卡通信開銷不是所有場景都值得。我在一些基準(zhǔn)測試?yán)镉^察過DDP開啟syncBN和不開最終精度差異基本在0.5%以內(nèi)但對訓(xùn)練吞吐影響不小。4. BN在網(wǎng)絡(luò)中放在哪里位置問題4.1 卷積后、激活前還是激活后這個問題的標(biāo)準(zhǔn)答案在業(yè)界經(jīng)歷了演化。原論文里BN是放在線性層之后、激活函數(shù)之前。在這個位置BN先對卷積輸出做標(biāo)準(zhǔn)化和縮放再送給ReLUReLU會截斷所有負(fù)值從而產(chǎn)生非對稱分布。這種做法的合理性在于卷積輸出的分布本身近似對稱歸一化后喂給ReLU能保證大多數(shù)激活值落在ReLU的有效區(qū)間內(nèi)。但在ResNet出現(xiàn)之前的一段時間也有不少研究工作把BN放在激活函數(shù)之后效果在某些任務(wù)上略好原因是ReLU后的分布大量零值經(jīng)過BN相當(dāng)于對稀疏激活做重新校準(zhǔn)。后來社區(qū)逐漸達(dá)成了共識對CNN而言把BN放在卷積層之后、激活函數(shù)之前是更通用的做法ResNet系列全部采用這個順序。我個人的經(jīng)驗也是如此這個位置在絕大多數(shù)視覺任務(wù)上表現(xiàn)最穩(wěn)。不過跟這個順序相關(guān)的一個小坑是BN對輸入分布的對稱性有隱含假設(shè)。如果輸入是ReLU輸出全是非負(fù)值其均值恒大于0BN為了歸一化就得強(qiáng)行把這個偏置分布搬回以0為中心這會壓縮原始激活的稀疏結(jié)構(gòu)有時反而有害。這也是為什么對激活后做BN一直沒成為主流的原因之一。4.2 全連接層里的BN怎么理解全連接層的BN跟卷積層的BN在原理上沒有區(qū)別但細(xì)節(jié)上有一個“通道”概念的變化。卷積層的BN是對每個輸出通道做統(tǒng)計而全連接層通常是每個神經(jīng)元一個維度嚴(yán)格來說BN是對每個輸出維度feature維度獨立做歸一化。實現(xiàn)上全連接層的BN層形狀是(num_features,)卷積層的BN形狀是(num_channels,)。實踐中的一個小經(jīng)驗全連接層對BN的收益往往不如卷積層明顯。因為全連接層處理的是“已經(jīng)語義化”的高層特征這些特征通常已經(jīng)被前面的卷積層很好地規(guī)整過了再強(qiáng)制歸一化帶來的收益有限有時甚至?xí)休p微的性能損失。所以我在設(shè)計網(wǎng)絡(luò)時更傾向于在卷積堆疊之后只在需要的地方加BN而不是給每個FC層都無腦加。4.3 在RNN/LSTM和Transformer里用BN為什么別扭把BN直接用在循環(huán)神經(jīng)網(wǎng)絡(luò)上效果一直不太好原因有幾個。第一RNN是時間步共享參數(shù)的同一個BN層在不同時間步上處理的數(shù)據(jù)分布其實一直在變。如果對序列的每個時間步都套用一個BN統(tǒng)計量就等于在給“不同分布”的數(shù)據(jù)用同一個標(biāo)準(zhǔn)化參數(shù)這本身就矛盾。有人說可以對每個時間步分別維護(hù)統(tǒng)計量但這樣參數(shù)數(shù)量隨序列長度增長且每步的樣本數(shù)往往很小統(tǒng)計量噪聲大到起不了作用。第二序列數(shù)據(jù)長度變化大時padding帶來的大量無效token會嚴(yán)重污染統(tǒng)計量。當(dāng)batch內(nèi)有不同長度的序列時短序列的padding位置對應(yīng)的激活值沒有實際意義把它們納入均值和方差的計算顯然是不合理的。第三是更根本的問題循環(huán)結(jié)構(gòu)天然要求“同一個處理邏輯在所有時間步上一致”而BN引入的batch級統(tǒng)計量把樣本間的信息混入到逐時間步的計算中破壞了這種一致性假設(shè)。所以在NLP、語音這類序列任務(wù)中更常見的選擇是LayerNorm層歸一化。這在Transformer架構(gòu)中幾乎是標(biāo)配BERT類模型里每個子層后面都跟著LayerNorm。LayerNorm跟BN最大的區(qū)別是它對一個樣本的所有特征維度做歸一化而不是對一個特征維度跨樣本做歸一化因此跟batch_size完全無關(guān)對變長序列也友好得多。5. 代碼實現(xiàn)從零手寫一個BN層5.1 為什么值得手寫一遍現(xiàn)在幾乎所有深度學(xué)習(xí)框架都內(nèi)置了BN調(diào)用就一兩行代碼。但我仍然強(qiáng)烈建議每個做深度學(xué)習(xí)的人至少手寫一遍BN的前向和反向過程。原因很簡單只有親手實現(xiàn)過才能真正理解訓(xùn)練和推理時統(tǒng)計量的流動、以及訓(xùn)練時“用的是當(dāng)前batch、維護(hù)的是全局狀態(tài)”這種雙重構(gòu)造的細(xì)節(jié)。在做框架遷移、自定義層、或者排查“訓(xùn)練好測試崩”這類問題的時候這些理解是直接能派上用場的。手寫B(tài)N不難關(guān)鍵是兩步前向傳播要用對訓(xùn)練/推理兩種模式反向傳播要能正確地把梯度分解回輸入。下面以PyTorch的自動梯度為基礎(chǔ)實現(xiàn)一個訓(xùn)練模式下的BN前向過程。5.2 一個教學(xué)版實現(xiàn)import torch import torch.nn as nn class ManualBatchNorm1d(nn.Module): def __init__(self, num_features, eps1e-5, momentum0.9): super().__init__() self.eps eps self.momentum momentum # 可學(xué)習(xí)參數(shù) self.gamma nn.Parameter(torch.ones(num_features)) self.beta nn.Parameter(torch.zeros(num_features)) # 全局統(tǒng)計量不參與梯度計算 self.register_buffer(running_mean, torch.zeros(num_features)) self.register_buffer(running_var, torch.ones(num_features)) def forward(self, x): # x: [batch_size, num_features] if self.training: # 1. 計算當(dāng)前batch的均值與方差有偏方差 batch_mean x.mean(dim0) batch_var x.var(dim0, unbiasedFalse) # 2. 歸一化 仿射變換 x_hat (x - batch_mean) / torch.sqrt(batch_var self.eps) out self.gamma * x_hat self.beta # 3. 更新全局統(tǒng)計量 # 注意PyTorch官方實現(xiàn)中跑的是無偏方差這里保持一致 unbiased_var x.var(dim0, unbiasedTrue) self.running_mean self.momentum * self.running_mean (1 - self.momentum) * batch_mean self.running_var self.momentum * self.running_var (1 - self.momentum) * unbiased_var else: # 推理模式使用全局統(tǒng)計量 x_hat (x - self.running_mean) / torch.sqrt(self.running_var self.eps) out self.gamma * x_hat self.beta return out這段代碼跟框架內(nèi)置BN在行為上已經(jīng)非常接近了。注意幾個關(guān)鍵細(xì)節(jié)x.mean(dim0)是對batch維做平均每個特征維度得到一個標(biāo)量unbiasedFalse表示用有偏方差歸一化更新running_var時換成了無偏方差是為了跟主流框架對齊這個細(xì)節(jié)第3章已經(jīng)講過。用這個自寫層做一次簡單驗證構(gòu)造一個隨機(jī)輸入對比自寫層跟框架內(nèi)置nn.BatchNorm1d的輸出在相同輸入、相同gamma/beta初始值下兩者的訓(xùn)練模式輸出應(yīng)該完全一致忽略浮點誤差推理模式在跑過幾個batch后也會趨于一致。5.3 真正需要手寫反向傳播嗎如果你的目的是理解原理那么借助自動求導(dǎo)跑通上述前向過程就足夠了因為梯度可以由框架自動計算。但如果你想徹底吃透BN我建議手動推導(dǎo)一次反向傳播的梯度公式并用torch.autograd.gradcheck來驗證手寫反向是否正確。這里只說推導(dǎo)的兩個關(guān)鍵難點。第一個難點是均值和方差的梯度必須“攤回”到每個樣本上。對$\hat{x}_i (x_i - \mu)/\sigma$求導(dǎo)時$\mu$和$\sigma$都依賴于所有$x$所以$\partial y_i/\partial x_j$在$i \neq j$時一般也不為0。這意味著反向傳播時每個輸入的梯度都包含了對同batch其他樣本的間接貢獻(xiàn)。手寫時必須先算$\partial y/\partial \hat{x}$、$\partial y/\partial \mu$、$\partial y/\partial \sigma^2$再組合出每個$x_i$的梯度。第二個難點是數(shù)值穩(wěn)定性。手動實現(xiàn)中如果直接使用$\sqrt{\sigma^2\epsilon}$作為分母當(dāng)方差很小時梯度會變得很大??蚣軆?nèi)置實現(xiàn)大多對此做了數(shù)值上的優(yōu)化比如用rsqrt結(jié)合torch.rsqrt或者重寫為不會產(chǎn)生極小分母的公式。自己做實驗時要注意同樣的問題不要太相信一個“裸寫”的BN在極小心方差下的梯度穩(wěn)定性。對大多數(shù)人而言到“能寫出正確前向、理解梯度流動方向”這個程度就夠了不必真的每層梯度都手寫出來。真正的價值在于弄懂那兩條統(tǒng)計量的路徑一條是當(dāng)前batch瞬時統(tǒng)計量的前向路徑另一條是running統(tǒng)計量的滑動更新路徑。6. 使用BN的訓(xùn)練技巧與調(diào)參經(jīng)驗6.1 學(xué)習(xí)率可以大膽放大的邊界在哪加了BN以后訓(xùn)練的一個直觀好處是學(xué)習(xí)率的可調(diào)整空間大了很多。沒有BN時我常用的初始學(xué)習(xí)率是0.01左右再往上就提心吊膽。有了BN在ResNet這類結(jié)構(gòu)上直接從0.1起步是常見做法甚至0.2也能訓(xùn)起來。但“能訓(xùn)起來”和“performance最好”是兩回事。學(xué)習(xí)率過大的典型癥狀是前期loss快速下降到中期開始震蕩最終收斂精度反而不如中等學(xué)習(xí)率。我通常的做法是有BN的結(jié)構(gòu)初始學(xué)習(xí)率設(shè)為無BN結(jié)構(gòu)能用的最大學(xué)習(xí)率的2到5倍然后用一個簡單的小數(shù)據(jù)集快速跑50個batch觀察loss是不是還在合理范圍內(nèi)波動以此快速試探學(xué)習(xí)的邊界。如果loss出現(xiàn)明顯的階梯式暴漲就說明學(xué)習(xí)率過大了。另一個跟BN相關(guān)的學(xué)習(xí)率經(jīng)驗是warmup。在超大batch訓(xùn)練以及Transformer類模型里warmup幾乎是必須的但對加了BN的CNN在batch規(guī)模正常比如32到256時warmup帶來的收益并不明顯可以直接跳過。只有在batch_size非常大比如1024以上時因為大batch讓統(tǒng)計量非常穩(wěn)定初始階段梯度方向一致但步長過大的風(fēng)險上升才需要配合warmup一起使用。6.2 batch size是BN最敏感的“外部旋鈕”BN對batch size的敏感程度是我在實際項目里感受最深的一點。它直接決定了訓(xùn)練時統(tǒng)計量的質(zhì)量batch越大均值和方差越穩(wěn)定訓(xùn)練越平滑batch越小統(tǒng)計量噪聲越大BN的行為就越不穩(wěn)定。我踩過一個印象很深的坑在某個目標(biāo)檢測項目里因為顯存限制batch_size只能設(shè)成2訓(xùn)練時loss曲線震蕩非常厲害模型勉強(qiáng)能收斂但mAP比預(yù)期低兩三個點。后來仔細(xì)排查發(fā)現(xiàn)BN統(tǒng)計量噪聲是主因。解決思路有三個方向第一個方向是用GroupNorm或LayerNorm替代BN。在檢測、分割這類天然需要大輸入分辨率、batch只能設(shè)很小的任務(wù)里這已經(jīng)不是“取巧”而是行業(yè)內(nèi)的標(biāo)準(zhǔn)做法。GroupNorm不需要跨樣本統(tǒng)計所以對batch大小完全免疫。第二個方向是盡量保持batch在32以上。如果卡顯存不夠就減小輸入分辨率、使用梯度累積保證BN的統(tǒng)計量有足夠樣本支撐。梯度累積雖然不直接降低統(tǒng)計量噪聲每個微batch還是很小但如果累積后再做歸一化是可以改變統(tǒng)計量計算范圍的——前提是你手動把統(tǒng)計量改為在整個累積batch上計算這實現(xiàn)起來有些麻煩多數(shù)人不會這么做。第三個方向是用SyncBN。在多卡環(huán)境下SyncBN會把所有卡上的數(shù)據(jù)合并起來計算全局統(tǒng)計量等效batch_size變大。我在8卡訓(xùn)練時開SyncBN后模型在小batch單卡場景下的表現(xiàn)明顯穩(wěn)定了。但要說清楚SyncBN不是銀彈它增加通信開銷單卡訓(xùn)練用不上它。6.3 BN跟Dropout、權(quán)重衰減、數(shù)據(jù)增強(qiáng)怎么配合BN自帶輕微正則化效果因此跟Dropout搭配時要重新審視Dropout的力度。我在圖像分類實驗里的經(jīng)驗是加了BN的模型用相同dropout比例訓(xùn)練損失和驗證損失的差距通常比無BN時小說明泛化差距已經(jīng)被BN“消化”了一部分。強(qiáng)行疊加高比例Dropout反而容易欠擬合驗證集精度不升反降?,F(xiàn)在主流的CNN結(jié)構(gòu)ResNet、DenseNet、EfficientNet都干脆不用Dropout至少不用在全連接之前的層上。權(quán)重衰減weight decay跟BN的關(guān)系就更微妙了。因為BN層里的γ和β也是可學(xué)習(xí)參數(shù)它們同樣會參與權(quán)重衰減。γ衰減過狠會限制BN的縮放表達(dá)能力β衰減過狠則會把偏移推向0。好在實踐中γ和β的權(quán)重衰減影響微乎其微絕大多數(shù)框架默認(rèn)對BN參數(shù)不做特殊的排除處理結(jié)果也沒出過毛病。不過在小數(shù)據(jù)集上過擬合嚴(yán)重時可以考慮把BN的γ和β從權(quán)重衰減中排除有時能帶來細(xì)微的精度提升。數(shù)據(jù)增強(qiáng)對BN的影響相對間接但值得注意當(dāng)數(shù)據(jù)增強(qiáng)強(qiáng)度很高時比如隨機(jī)裁剪加色彩抖動每個batch內(nèi)的樣本多樣性變大BN統(tǒng)計量的方差也隨之增大。這其實不是壞事因為更“多樣”的統(tǒng)計量反而起到了更強(qiáng)的正則化作用。但如果增強(qiáng)過于激進(jìn)導(dǎo)致batch內(nèi)出現(xiàn)明顯的離群點均值和方差容易被少數(shù)極端樣本拉偏個別情況下會出現(xiàn)訓(xùn)練不穩(wěn)定。我在AutoAugment強(qiáng)度翻倍后再配合小batch訓(xùn)練時確實遇到過loss陡增的情況后來把batch_size調(diào)大后問題自然消失了。7. 常見問題與排查實錄7.1 訓(xùn)練正常但測試崩了先查狀態(tài)切換這種“訓(xùn)練loss很好看測試輸出一團(tuán)糟”的案例我大概見過幾十次。絕大多數(shù)時候原因驚人地一致某個環(huán)節(jié)忘了切model.eval()或者跑推理的腳本里壓根沒有調(diào)用。這類bug最麻煩的地方在于在batch較大的推理場景下用當(dāng)前batch統(tǒng)計量歸一化并不一定會讓結(jié)果“肉眼可見地差”。很多時候模型輸出甚至看起來合理只有當(dāng)你用單樣本推理、或者對比相同輸入在不同batch下的輸出時才暴露問題——輸出不一致。排查方法很簡單寫個測試固定輸入張量A分別以batch1和batch8跑前向比較輸出。如果輸出有明顯差異幾乎可以斷定模型處于training模式。再檢查一下是否所有子模塊的BN都被正確地設(shè)為eval有時模型里某個自定義層內(nèi)部偷偷把BN的training改了也會導(dǎo)致這類問題。7.2 小batch訓(xùn)練不收斂別硬調(diào)學(xué)習(xí)率當(dāng)你發(fā)現(xiàn)batch在2到8之間訓(xùn)練loss怎么調(diào)都不對勁時先別急著懷疑學(xué)習(xí)率、優(yōu)化器、網(wǎng)絡(luò)結(jié)構(gòu)大概率就是BN統(tǒng)計量樣本太少導(dǎo)致的。判斷方法可以這樣把batch增大到32或64如果可以的話如果同樣代碼下loss曲線立刻變得平滑且收斂正常那基本確認(rèn)是BN統(tǒng)計量噪聲的問題。這種時候我建議直接在網(wǎng)絡(luò)里把關(guān)鍵層的BN換成GroupNorm或LayerNorm而不是嘗試通過調(diào)低學(xué)習(xí)率來“硬扛”。為什么換層比調(diào)學(xué)習(xí)率更有效因為小batch下BN的歸一化本身就不穩(wěn)定導(dǎo)致梯度更新信號被污染學(xué)習(xí)率再低也只是緩解癥狀沒有根治。換用不依賴batch的歸一化方式相當(dāng)于把這個不穩(wěn)定源從根上消滅。我在一個語義分割項目里把整條backbone的BN換成GN后batch從4降到2精度幾乎無損訓(xùn)練曲線也平穩(wěn)了很多這是實測過最有效的方案。7.3 loss曲線異常震蕩時怎么快速定位是不是BN的鍋loss震蕩的原因很多不一定全是BN的問題。一個快速定位方法訓(xùn)練前把BN的momentum暫時加大比如0.99改成0.5讓統(tǒng)計量快速跟蹤當(dāng)前batch的變化同時把γ初始化改成0這樣BN輸出恒為β歸一化分支暫時不起作用看loss的表現(xiàn)。如果人為抑制BN后loss反而穩(wěn)定說明問題與BN高度相關(guān)如果不相關(guān)問題更可能在數(shù)據(jù)管線、學(xué)習(xí)率或模型架構(gòu)上。這里補(bǔ)充一個工程細(xì)節(jié)momentum調(diào)大只是權(quán)宜之計不要把它當(dāng)長期訓(xùn)練的解藥。momentum過大太接近0會讓running統(tǒng)計量頻繁大幅波動不僅訓(xùn)練后期統(tǒng)計量不穩(wěn)定還會讓最終模型的推理性能變差因為訓(xùn)練和推理時的統(tǒng)計量差異越來越大。7.4 微調(diào)預(yù)訓(xùn)練模型時BN到底要不要凍結(jié)遷移學(xué)習(xí)場景下關(guān)于BN的爭議一直存在微調(diào)時是凍結(jié)BN層還是放開。我在多個項目里做過對照實驗結(jié)論是如果是小數(shù)據(jù)集幾百到幾千張微調(diào)凍結(jié)BN層通常效果更好如果是大數(shù)據(jù)集幾萬張以上或完整微調(diào)放開BN通常能帶來更高上限。原因也好理解預(yù)訓(xùn)練模型的BN統(tǒng)計量是在大規(guī)模、多樣本上積累出來的小數(shù)據(jù)集里每個batch的統(tǒng)計量噪聲大直接放開BN等于讓一個“訓(xùn)練充分”的統(tǒng)計量被少量樣本迅速帶偏反而破壞預(yù)訓(xùn)練特征的分布。凍結(jié)BN相當(dāng)于保留預(yù)訓(xùn)練階段的數(shù)據(jù)分布先驗?zāi)P椭煌ㄟ^調(diào)整其他層的參數(shù)來適應(yīng)新任務(wù)。實踐中我還遇到過另一個極端用某個預(yù)訓(xùn)練模型做超小數(shù)據(jù)集100張圖的二分類放開BN訓(xùn)練后驗證集精度從90%掉到70%以下。凍結(jié)BN后回到90%。這個案例很典型建議微調(diào)新手先把“凍結(jié)BN層”作為一個默認(rèn)選項再根據(jù)實驗結(jié)果決定是否放開。7.5 一個快速問題排查速查表癥狀常見原因排查方向訓(xùn)練正常單樣本測試輸出異常模型未切evaltrack_running_statsFalse檢查model.eval()檢查BN參數(shù)配置小batch訓(xùn)練不收斂BN統(tǒng)計量噪聲過大增大batch或替換為GroupNorm/LayerNorm多卡訓(xùn)練精度低于單卡BN統(tǒng)計量未同步或同步時機(jī)不對檢查DDP的syncBN配置微調(diào)預(yù)訓(xùn)練模型掉點BN統(tǒng)計量被小數(shù)據(jù)集污染凍結(jié)BN層加載模型后測試輸出全亂running_mean/running_var缺失檢查state_dict是否包含BN統(tǒng)計量訓(xùn)練后期loss周期性跳動統(tǒng)計量更新不穩(wěn)定可能momentum過小恢復(fù)momentum默認(rèn)值0.9~0.998. BN的局限與后續(xù)變體什么時候該換掉BN8.1 BN的兩個結(jié)構(gòu)性“硬傷”BN最被詬病的兩個問題都跟“batch”二字綁定。一個是batch大小波動的脆弱性訓(xùn)練時batch很大、上線后如果為了延遲強(qiáng)行降batch遇到的效果會跟訓(xùn)練時完全不同。另一個是訓(xùn)練和推理行為的割裂訓(xùn)練時模型看到的歸一化統(tǒng)計量是動態(tài)的、有噪聲的推理時卻換成了固定統(tǒng)計量這個gap在小batch、大momentum、分布漂移三種因素疊加時會變得非常明顯。這兩個“硬傷”不是靠調(diào)參能消除的本質(zhì)上是BN定義上對batch維度的依賴造成的。所以當(dāng)你的任務(wù)場景涉及“在線推理必須低延遲”“batch被迫很小”“序列樣本長度變化劇烈”三者之一時認(rèn)真考慮替代方案是值得的。8.2 GroupNorm、LayerNorm、InstanceNorm的選擇邏輯這幾類歸一化方法本質(zhì)區(qū)別在于沿著哪幾個維度做統(tǒng)計。用圖像特征圖N, C, H, W來表述最清晰BatchNorm對(N, H, W)做統(tǒng)計保留C維度。適合batch充足的視覺任務(wù)。LayerNorm對(C, H, W)做統(tǒng)計保留N維度。適合NLP、Transformer對batch不敏感。InstanceNorm對(H, W)做統(tǒng)計保留N和C。適合風(fēng)格遷移等需要對單張圖像做歸一化的任務(wù)。GroupNorm對(C分成G組每組里的C/G, H, W)做統(tǒng)計保留N和組維度。是batch受限時CNNs的主流替代。方法統(tǒng)計維度依賴batch適用場景BatchNormN,H,W是視覺任務(wù)、batch≥32LayerNormC,H,W否NLP、Transformer、變長序列InstanceNormH,W否風(fēng)格遷移、圖像生成GroupNorm分組C,H,W否檢測/分割等batch受限任務(wù)GroupNorm在目標(biāo)檢測、語義分割這類“輸入分辨率大、batch只能放幾張圖”的任務(wù)里已經(jīng)成為事實標(biāo)準(zhǔn)。我在多個檢測模型里做過GroupNorm與BN的轉(zhuǎn)換測試在等價訓(xùn)練配置下兩者精度基本持平但GroupNorm對batch變化的魯棒性明顯更好。8.3 對BN本質(zhì)的再思考跟BN打了這么多年交道我的體會是它本質(zhì)上是在“特征分布控制”和“樣本間信息共享”之間尋找平衡。它控制分布的能力非常強(qiáng)代價是引入了對batch的依賴和樣本間的耦合。理解這一點你就不會在面對“為什么不直接在所有地方都用BN”的問題時犯迷糊——不是BN不好而是它天生帶有一個batch大小的“適用范圍”。從工程角度看一個好的策略不是迷信某一種歸一化而是根據(jù)任務(wù)形態(tài)選型主流CNN 充足batch用BNbatch受限的視覺任務(wù)用GroupNorm序列模型用LayerNorm。這幾種方案之間的轉(zhuǎn)換成本很低通常只需要替換網(wǎng)絡(luò)中的歸一化層其他結(jié)構(gòu)不用動。但一定要在項目早期就做這個決策因為不同歸一化對超參數(shù)學(xué)習(xí)率、weight decay、batch size的敏感度不一樣中途更換意味著要把訓(xùn)練配置也重新調(diào)一遍這個成本往往被低估。在我個人的實踐里BN仍然是我接觸一個新視覺任務(wù)時的默認(rèn)起點因為它的收斂速度和穩(wěn)定表現(xiàn)確實讓人省心。但我也學(xué)會了時刻問自己一個問題當(dāng)前任務(wù)的batch size真的足夠支撐BN的統(tǒng)計量嗎如果答案不確定直接用GroupNorm可能省下好幾天的調(diào)參時間。畢竟做工程項目的目標(biāo)是把模型又快又好地跑起來而不是跟某個特定的歸一化方法“死磕到底”。