實(shí)戰(zhàn)：凍結(jié)層參數(shù)的原理、策略與踩坑指南)
不少朋友第一次接觸模型微調(diào)時(shí)都會(huì)遇到一個(gè)困惑為什么大家都在說(shuō)“凍結(jié)網(wǎng)絡(luò)”、“freeze層參數(shù)”這個(gè)詞聽(tīng)起來(lái)很玄乎其實(shí)說(shuō)白了就是讓某些層的權(quán)重在訓(xùn)練過(guò)程中保持不變只更新剩下的部分。PyTorch里做這件事非常靈活可以用requires_grad False控制梯度也可以在優(yōu)化器里手動(dòng)過(guò)濾參數(shù)還可以配合eval()模式處理BatchNorm這類有狀態(tài)層的坑。這篇文章我想把這套東西徹底講明白。我會(huì)從為什么要凍結(jié)參數(shù)講起然后展示幾種常見(jiàn)的凍結(jié)姿勢(shì)再結(jié)合一個(gè)貼近真實(shí)業(yè)務(wù)的多任務(wù)模型例子把“凍結(jié)哪些層、怎么解凍、學(xué)習(xí)率怎么配”這些實(shí)操細(xì)節(jié)全部過(guò)一遍。你可能是剛?cè)腴TPyTorch的新手也可能已經(jīng)在做遷移學(xué)習(xí)但總被BN層坑得頭皮發(fā)麻這篇文章的思路和踩坑經(jīng)驗(yàn)對(duì)你應(yīng)該都有直接幫助。1. 為什么要凍結(jié)網(wǎng)絡(luò)參數(shù)從省顯存到防過(guò)擬合先說(shuō)結(jié)論凍結(jié)參數(shù)這個(gè)操作本質(zhì)上是在干預(yù)優(yōu)化過(guò)程。神經(jīng)網(wǎng)絡(luò)靠反向傳播更新參數(shù)每一步都需要計(jì)算損失函數(shù)對(duì)各層權(quán)重的梯度。如果把某些層的requires_grad設(shè)為False梯度根本不會(huì)流到那里參數(shù)自然保持不動(dòng)。這一招最常見(jiàn)的應(yīng)用場(chǎng)景有三個(gè)。第一個(gè)場(chǎng)景是遷移學(xué)習(xí)。你拿一個(gè)在大規(guī)模數(shù)據(jù)集上預(yù)訓(xùn)練好的模型當(dāng)特征提取器比如圖像分類里的ResNet、文本分類里的BERT這些模型底層已經(jīng)學(xué)會(huì)了通用的視覺(jué)紋理或句法結(jié)構(gòu)。下游任務(wù)的訓(xùn)練數(shù)據(jù)往往只有幾千張圖、幾萬(wàn)條文本如果從零開(kāi)始訓(xùn)練所有層模型很容易把訓(xùn)練集里的噪聲背下來(lái)。把預(yù)訓(xùn)練部分的參數(shù)凍住只訓(xùn)練最后新增的分類頭相當(dāng)于站在巨人肩膀上做微調(diào)。這個(gè)思路在數(shù)據(jù)量小、任務(wù)與預(yù)訓(xùn)練任務(wù)相似度高的場(chǎng)景下尤其好用。第二個(gè)場(chǎng)景是顯存和算力限制。訓(xùn)練一個(gè)大規(guī)模模型顯存消耗主要來(lái)自中間激活值和優(yōu)化器狀態(tài)。凍結(jié)部分層之后PyTorch在前向傳播時(shí)依然需要算這些層的輸出但反向傳播不再為它們緩存激活值優(yōu)化器也不再為這些層保存動(dòng)量或二階矩估計(jì)。我實(shí)測(cè)過(guò)一個(gè)約3億參數(shù)量的模型凍結(jié)掉底部80%的層之后訓(xùn)練峰值顯存直接降了將近一半單卡跑批大小也能翻倍。第三個(gè)場(chǎng)景容易被人忽略——防止災(zāi)難性遺忘。在做增量學(xué)習(xí)或者多任務(wù)學(xué)習(xí)時(shí)如果希望模型在完成新任務(wù)的同時(shí)不丟掉舊任務(wù)的能力凍結(jié)掉共享底層參數(shù)是性價(jià)比最高的方案之一。我在某個(gè)多語(yǔ)言文本分類項(xiàng)目中一開(kāi)始讓所有層都參與訓(xùn)練結(jié)果發(fā)現(xiàn)模型在新語(yǔ)言上的指標(biāo)上去了舊語(yǔ)言的F1值卻掉得厲害。后來(lái)把共享編碼器的前幾層凍住只訓(xùn)練頂層和任務(wù)專屬的分類分支舊任務(wù)的表現(xiàn)就穩(wěn)定下來(lái)了。當(dāng)然凍結(jié)參數(shù)不是萬(wàn)能的它也有明顯的適用邊界。比如預(yù)訓(xùn)練模型和下游任務(wù)差異巨大時(shí)底層特征完全不更新會(huì)導(dǎo)致模型學(xué)不到任務(wù)相關(guān)的特定表示。再比如你對(duì)某個(gè)任務(wù)的標(biāo)注數(shù)據(jù)足夠多那讓所有層都參與微調(diào)通常能拿到更好的最終精度。所以實(shí)際操作中我一般不會(huì)一上來(lái)就把模型焊死而是先做幾組“全量微調(diào) vs 部分凍結(jié)”的對(duì)照實(shí)驗(yàn)看指標(biāo)差距再?zèng)Q定。2. 核心機(jī)制拆解requires_grad、優(yōu)化器過(guò)濾與訓(xùn)練/評(píng)估模式要把凍結(jié)參數(shù)玩明白先得弄懂PyTorch底層那套梯度傳播邏輯。很多人只知道設(shè)requires_gradFalse就行但對(duì)它背后的機(jī)制、和eval()的搭配、還有優(yōu)化器行為之間的關(guān)系并不清楚結(jié)果總是遇到“明明凍住了參數(shù)怎么還在變”這種怪事。2.1requires_gradFalse的傳播阻斷原理PyTorch的自動(dòng)求導(dǎo)引擎基于計(jì)算圖工作。每當(dāng)你做一次前向傳播框架會(huì)記錄哪些張量參與了運(yùn)算反向傳播時(shí)沿著這個(gè)圖把梯度從損失函數(shù)一路傳回各層的權(quán)重矩陣。requires_gradFalse的作用非常直接這個(gè)張量不會(huì)再出現(xiàn)在計(jì)算圖的葉子節(jié)點(diǎn)列表里反向傳播經(jīng)過(guò)它時(shí)直接跳過(guò)梯度傳不到它身上參數(shù)自然紋絲不動(dòng)。這里有個(gè)細(xì)節(jié)值得說(shuō)透——即使你凍結(jié)了某一層前向傳播仍然會(huì)正常計(jì)算該層的輸出結(jié)果這個(gè)結(jié)果會(huì)被繼續(xù)送入后續(xù)層。也就是說(shuō)網(wǎng)絡(luò)結(jié)構(gòu)沒(méi)有變“前向推理路徑”是完整的只是“參數(shù)的更新路徑”被切斷了。我在給別人講這個(gè)概念時(shí)經(jīng)常打一個(gè)比方凍結(jié)參數(shù)相當(dāng)于把一臺(tái)機(jī)器里某個(gè)調(diào)節(jié)旋鈕擰緊并鎖死機(jī)器還在運(yùn)轉(zhuǎn)產(chǎn)品還是會(huì)出來(lái)只是這個(gè)旋鈕以后不會(huì)再被工人調(diào)整了。寫代碼時(shí)要注意requires_grad必須遞歸設(shè)置到所有子模塊上。比如你用了torchvision.models.resnet50想凍結(jié)整個(gè)特征提取器只寫model.layer1.requires_grad_(False)可不夠因?yàn)閘ayer1里面還有若干個(gè)Bottleneck模塊每個(gè)模塊里面又有多個(gè)卷積層和BN層。正確的做法是遍歷模塊遞歸設(shè)置def freeze_module(module): for param in module.parameters(): param.requires_grad False這個(gè)函數(shù)會(huì)遍歷傳入模塊的所有子模塊并凍結(jié)其參數(shù)比直接操作某個(gè)模塊要穩(wěn)妥得多。2.2 優(yōu)化器里的“隱形坑”參數(shù)過(guò)濾與緩存狀態(tài)很多人在凍結(jié)參數(shù)后忽略了一個(gè)關(guān)鍵步驟——優(yōu)化器需要重新構(gòu)建。假設(shè)你先創(chuàng)建了一個(gè)優(yōu)化器傳入所有參數(shù)然后才把某些層設(shè)成requires_gradFalse那么優(yōu)化器內(nèi)部依然保留著這些參數(shù)的引用并且會(huì)在每個(gè)step里試圖更新它們。雖然由于梯度為0參數(shù)實(shí)際上不會(huì)變化但這會(huì)白白浪費(fèi)計(jì)算和顯存而且沒(méi)有意義的動(dòng)量狀態(tài)累積起來(lái)還可能留下隱患。正確做法是先用requires_gradFalse凍結(jié)參數(shù)再創(chuàng)建優(yōu)化器并且只把那些仍然requires_gradTrue的參數(shù)傳進(jìn)去frozen_params [] trainable_params [] for name, param in model.named_parameters(): if param.requires_grad: trainable_params.append(param) else: frozen_params.append(param) optimizer optim.Adam(trainable_params, lr1e-4)這里我用named_parameters()而不是parameters()目的是拿到每個(gè)參數(shù)的名字方便后續(xù)打印、篩選或者按層名稱做精細(xì)化的學(xué)習(xí)率分組。而且把這個(gè)參數(shù)列表打印出來(lái)你會(huì)很清楚當(dāng)前到底哪些層在訓(xùn)練、哪些層是凍結(jié)的。另一個(gè)容易踩的坑是優(yōu)化器緩存。某些優(yōu)化器比如SGD的動(dòng)量、Adam的一階二階動(dòng)量是根據(jù)參數(shù)狀態(tài)來(lái)維護(hù)緩存的。你雖然已經(jīng)設(shè)了requires_gradFalse但舊優(yōu)化器里關(guān)于這些參數(shù)的緩存可能依然存在。更穩(wěn)妥的做法是只要你調(diào)整了凍結(jié)范圍就直接重新創(chuàng)建一個(gè)優(yōu)化器實(shí)例不要復(fù)用舊的。這不用寫多少代碼但能有效避免很多詭異的問(wèn)題。2.3train()與eval()模式不只是Dropout的開(kāi)關(guān)初學(xué)者很容易把參數(shù)凍結(jié)和model.eval()混為一談。它倆的確有交集但完全是兩碼事。model.eval()會(huì)把模型切到評(píng)估模式影響的是Dropout、BatchNorm這類在訓(xùn)練和推理時(shí)行為不同的層但它不會(huì)把任何參數(shù)設(shè)成requires_gradFalse。也就是說(shuō)即使你調(diào)用了model.eval()只要優(yōu)化器里還有這些參數(shù)它們照樣會(huì)被更新。反過(guò)來(lái)凍結(jié)參數(shù)也不意味著你要把整個(gè)模型切到eval()模式。這里有個(gè)典型誤區(qū)很多人凍結(jié)了特征提取層之后習(xí)慣性地調(diào)用了model.eval()結(jié)果BN層在推理時(shí)用的是全局統(tǒng)計(jì)量而不是當(dāng)前batch的統(tǒng)計(jì)量這反而影響了訓(xùn)練效果。更常見(jiàn)的做法是只對(duì)含BatchNorm的部分單獨(dú)處理或者干脆將BN層也放在訓(xùn)練模式下但凍結(jié)其中的可學(xué)習(xí)參數(shù)。在動(dòng)手之前請(qǐng)務(wù)必弄清楚這幾件事的邊界requires_gradFalse控制的是梯度計(jì)算和參數(shù)更新optimizer決定哪些參數(shù)參與更新model.train()/model.eval()控制的是層的行為模式。參數(shù)凍結(jié)時(shí)我一般會(huì)把最終要訓(xùn)練的那部分模塊保持train()模式但會(huì)根據(jù)BN層是否會(huì)被凍結(jié)來(lái)單獨(dú)處理。3. 實(shí)操?gòu)囊粋€(gè)多任務(wù)模型看不同凍結(jié)策略的代碼落地理論說(shuō)得再多不如看一個(gè)完整的實(shí)戰(zhàn)例子。我在很多業(yè)務(wù)場(chǎng)景中會(huì)遇到這樣的情況模型需要同時(shí)完成一個(gè)主任務(wù)和一個(gè)輔助任務(wù)但輔助任務(wù)的標(biāo)注數(shù)據(jù)非常有限如果讓所有參數(shù)都參與優(yōu)化很容易在輔助任務(wù)上過(guò)擬合。下面我用一個(gè)簡(jiǎn)化版本的多任務(wù)分類模型來(lái)演示不同的凍結(jié)策略。3.1 搭建一個(gè)共享底層、多分支輸出的模型假設(shè)我們有一個(gè)眾包標(biāo)注平臺(tái)需要同時(shí)預(yù)測(cè)圖片的“主類別”和“質(zhì)量分?jǐn)?shù)”。主類別的標(biāo)注數(shù)據(jù)很充裕質(zhì)量分?jǐn)?shù)的標(biāo)注卻比較少。一個(gè)合理的模型結(jié)構(gòu)是共享一個(gè)ResNet風(fēng)格的編碼器然后分出兩個(gè)分類頭。這個(gè)模擬項(xiàng)目X的目標(biāo)就是驗(yàn)證“凍結(jié)共享編碼器、只訓(xùn)練兩個(gè)分類頭”能否在輔助任務(wù)上取得更好的泛化效果。import torch import torch.nn as nn from torchvision import models class MultiTaskModel(nn.Module): def __init__(self): super().__init__() self.backbone models.resnet18(pretrainedTrue) self.backbone.fc nn.Identity() # 去掉原來(lái)的全連接層 self.head_main nn.Sequential( nn.Linear(512, 128), nn.ReLU(inplaceTrue), nn.Linear(128, 10) ) self.head_aux nn.Sequential( nn.Linear(512, 64), nn.ReLU(inplaceTrue), nn.Linear(64, 1) ) def forward(self, x): feat self.backbone(x) out_main self.head_main(feat) out_aux self.head_aux(feat) return out_main, out_aux這個(gè)模型很典型共享編碼器是預(yù)訓(xùn)練好的ResNet18兩個(gè)任務(wù)各有一個(gè)小分類頭。現(xiàn)在我們要把策略落在代碼上。3.2 策略一只訓(xùn)練分類頭凍結(jié)全部backbone參數(shù)這個(gè)策略最簡(jiǎn)單也是我早期最常用的方法。整個(gè)預(yù)訓(xùn)練編碼器一筆不動(dòng)只訓(xùn)練兩個(gè)新增的全連接分支。這樣做的好處是訓(xùn)練速度極快baseline跑起來(lái)特別快幾分鐘就能看到路線是否可行。def freeze_encoder_and_train_heads(model): # 凍結(jié)backbone for name, param in model.backbone.named_parameters(): param.requires_grad False # 只把head里的參數(shù)傳給優(yōu)化器 trainable_params [] for name, param in model.named_parameters(): if param.requires_grad: trainable_params.append(param) optimizer torch.optim.Adam(trainable_params, lr1e-3) return optimizer這里有個(gè)很容易被忽略的細(xì)節(jié)我們用的是model.named_parameters()但param.requires_grad是直接讀的參數(shù)屬性。如果某些層之前已經(jīng)被設(shè)置為不可訓(xùn)練這里會(huì)自動(dòng)跳過(guò)它們無(wú)需額外過(guò)濾。訓(xùn)練時(shí)你可以順手打印當(dāng)前正在訓(xùn)練的層名確認(rèn)自己的意圖真正落到了代碼層面for name, param in model.named_parameters(): if param.requires_grad: print(訓(xùn)練中的參數(shù):, name)3.3 策略二凍結(jié)backbone的前幾層解凍后面高層并設(shè)置差異化學(xué)習(xí)率凍結(jié)全部backbone最省事但上限也不高。如果任務(wù)和預(yù)訓(xùn)練任務(wù)差距比較大底層特征完全不動(dòng)的話最終精度往往不夠理想。實(shí)踐中我會(huì)用一個(gè)更精細(xì)的策略底層的淺層特征比如邊緣、顏色、紋理這些通用結(jié)構(gòu)完全凍結(jié)高層的語(yǔ)義特征比如跟具體物體部件相關(guān)的特征解凍微調(diào)兩個(gè)分類頭用更高的學(xué)習(xí)率backbone的高層用較低學(xué)習(xí)率。def partial_freeze_with_lr(model): # 凍結(jié)layer1和layer2解凍layer3和layer4 for name, param in model.backbone.named_parameters(): if name.startswith(layer1) or name.startswith(layer2): param.requires_grad False else: param.requires_grad True optimizer torch.optim.Adam([ # 分類頭大學(xué)習(xí)率 {params: model.head_main.parameters(), lr: 1e-3}, {params: model.head_aux.parameters(), lr: 1e-3}, # backbone高層小學(xué)習(xí)率 {params: filter(lambda p: p.requires_grad, model.backbone.layer3.parameters()), lr: 1e-5}, {params: filter(lambda p: p.requires_grad, model.backbone.layer4.parameters()), lr: 1e-5}, ]) return optimizer差異化學(xué)習(xí)率是這策略的靈魂。底層特征通用性強(qiáng)、更新風(fēng)險(xiǎn)高就完全鎖死高層特征與具體任務(wù)相關(guān)性強(qiáng)可以適當(dāng)更新但預(yù)訓(xùn)練的知識(shí)也別丟得太多所以用一個(gè)很小的學(xué)習(xí)率。分類頭是隨機(jī)初始化的要從頭學(xué)所以給它最大的學(xué)習(xí)率。這樣一組參數(shù)在分類準(zhǔn)確率和訓(xùn)練穩(wěn)定性上都比較平衡。關(guān)于解凍的邊界一個(gè)穩(wěn)妥的方法是觀察每層特征圖的激活變化。你可以拿一個(gè)固定batch的數(shù)據(jù)在前向傳播時(shí)打印每層輸出的均值方差對(duì)比凍結(jié)前后哪些層的輸出變化最大。如果發(fā)現(xiàn)layer2解凍后輸出分布已經(jīng)亂得不成樣子說(shuō)明學(xué)習(xí)率太高了或者干脆把它凍回去更穩(wěn)。3.4 策略三在凍結(jié)狀態(tài)下做“冷啟動(dòng)”再用“解凍全網(wǎng)微調(diào)”做熱階段有些任務(wù)雖然數(shù)據(jù)量一般但遷移學(xué)習(xí)的起點(diǎn)非常關(guān)鍵。我慣用的一個(gè)方法是兩階段訓(xùn)練第一階段讓backbone全面凍結(jié)只訓(xùn)練分類頭目標(biāo)是讓分類頭在預(yù)訓(xùn)練特征上找到一個(gè)大致的解第二階段解凍全部參數(shù)用一個(gè)非常小的學(xué)習(xí)率做全網(wǎng)微調(diào)。這比直接從隨機(jī)初始化訓(xùn)練整個(gè)網(wǎng)絡(luò)要穩(wěn)定不少也能明顯提升最終精度。# 階段一只訓(xùn)練head def stage_one(model): for p in model.backbone.parameters(): p.requires_grad False return torch.optim.Adam(model.head_main.parameters(), lr1e-3) # 階段二全部解凍 低學(xué)習(xí)率 def stage_two(model): for p in model.backbone.parameters(): p.requires_grad True optimizer torch.optim.Adam(model.parameters(), lr1e-5) return optimizer這里的關(guān)鍵點(diǎn)是階段一和階段二之間除了學(xué)習(xí)率變化最好把階段一的優(yōu)化器丟棄掉重新建一個(gè)新的。原因前面提過(guò)優(yōu)化器緩存里還保存著階段一的狀態(tài)突然換到全網(wǎng)微調(diào)時(shí)那些狀態(tài)可能會(huì)干擾更新方向重新創(chuàng)建更干凈。4. BN層、參數(shù)加載與多卡訓(xùn)練的連環(huán)坑凍結(jié)參數(shù)看起來(lái)是個(gè)小操作但在真實(shí)項(xiàng)目里總會(huì)碰到各種“看著沒(méi)事、跑起來(lái)就掛”的情況。這里面的坑我基本都踩過(guò)一遍一個(gè)個(gè)說(shuō)。4.1 BatchNorm層只凍結(jié)可學(xué)習(xí)參數(shù)遠(yuǎn)遠(yuǎn)不夠這是最容易踩、也最隱蔽的坑。BatchNorm層雖然是通過(guò)weight和bias學(xué)習(xí)縮放和平移但它的行為還依賴兩個(gè)“狀態(tài)量”——running_mean和running_var。這兩個(gè)量不是靠反向傳播更新的而是在每次前向傳播時(shí)用一個(gè)動(dòng)量系數(shù)在線更新的。所以即使你設(shè)了requires_gradFalse只要model.train()模式下這個(gè)BN層還在跑前向計(jì)算running_mean和running_var就會(huì)不斷被當(dāng)前batch的統(tǒng)計(jì)量更新。這在訓(xùn)練和推理時(shí)會(huì)出現(xiàn)嚴(yán)重的分布漂移問(wèn)題。如何正確凍結(jié)BatchNorm層我給你一個(gè)標(biāo)準(zhǔn)做法def freeze_bn(mod): for m in mod.modules(): if isinstance(m, nn.BatchNorm2d): m.eval() # 讓BN層使用全局統(tǒng)計(jì)量而不是更新running_mean for p in m.parameters(): p.requires_grad False這樣做的本質(zhì)是讓BN層切到評(píng)測(cè)模式前向傳播時(shí)直接使用running_mean和running_var不再用本batch的統(tǒng)計(jì)量更新它們。這樣才能做到真正的“凍結(jié)”。很多人只設(shè)了requires_gradFalse卻忘了讓BN層切模式結(jié)果模型在訓(xùn)練時(shí)表現(xiàn)正常、在驗(yàn)證時(shí)指標(biāo)掉得厲害八成就是這里出了問(wèn)題。那到底要凍結(jié)BN嗎需要分場(chǎng)景討論。如果凍結(jié)的是整個(gè)backboneBN層務(wù)必按上面這種方式處理因?yàn)樗鼉?nèi)部的統(tǒng)計(jì)數(shù)據(jù)如果持續(xù)變化會(huì)干擾凍結(jié)層輸出的特征一致性如果模型很小、任務(wù)分布和預(yù)訓(xùn)練分布差別很大BN也可以保持默認(rèn)的跟蹤更新只是這時(shí)它沒(méi)有被真正凍結(jié)你需要心里有數(shù)。4.2 加載預(yù)訓(xùn)練權(quán)重時(shí)嚴(yán)格對(duì)齊參數(shù)名和形狀用別人訓(xùn)練好的權(quán)重來(lái)微調(diào)如果你的模型結(jié)構(gòu)是從torchvision.models.resnet18(pretrainedTrue)直接加載的那名字天然對(duì)得上。但如果你自己改了結(jié)構(gòu)比如把fc換成一個(gè)雙分支頭加載權(quán)重時(shí)就很可能報(bào)size mismatch的錯(cuò)。這類報(bào)錯(cuò)原因很直接預(yù)訓(xùn)練權(quán)重的某個(gè)層形狀和你的模型對(duì)不上。比如原始的ResNet18最后一個(gè)全連接層是(512, 1000)形狀你把它換成了(512, 10)那加載時(shí)就肯定會(huì)沖突。解決辦法是strictFalse只加載能對(duì)上的部分checkpoint torch.load(pretrained_weights.pth) model.load_state_dict(checkpoint, strictFalse)但這里要小心strictFalse會(huì)把無(wú)法匹配的層全部忽略包括你本來(lái)想加載的層。要確保想加載的特征提取層能被正確覆蓋我建議加載前先打印一下檢查一下model.load_state_dict(checkpoint, strictFalse) for name, param in model.named_parameters(): print(name, param.shape, param.requires_grad)另外如果你是自己動(dòng)手做的模型加載之前還應(yīng)該確認(rèn)state_dict里的鍵名是否跟模型里的一致。比如某層以前叫backbone.conv1.weight后來(lái)你重構(gòu)成了encoder.conv1.weight鍵名就對(duì)不上了得手工做一下映射。最省事的做法是在模型保存和加載時(shí)都使用同一套命名規(guī)則別在中間改名。4.3 多卡訓(xùn)練時(shí)凍結(jié)邏輯與All-Reduce的交互使用DataParallel或DistributedDataParallel時(shí)凍結(jié)參數(shù)還有一個(gè)需要提前想清楚的地方?？蚣茉诿看畏聪騻鞑r(shí)會(huì)對(duì)梯度做跨卡同步All-Reduce。如果你只在一張卡上凍結(jié)參數(shù)而在其他卡上沒(méi)有保持一致梯度同步就會(huì)出問(wèn)題最終導(dǎo)致更新行為完全不可控。而且要特別注意在DistributedDataParallel模式下歸一化層和同步層本來(lái)就是全局同步的。如果通過(guò)某些方式把BN層也帶上了同步邏輯你凍結(jié)時(shí)也要考慮sync BN的行為。我的經(jīng)驗(yàn)是在做多卡訓(xùn)練前先把凍結(jié)參數(shù)的全部邏輯固定下來(lái)讓所有進(jìn)程在完全一致的參數(shù)狀態(tài)下進(jìn)入訓(xùn)練在訓(xùn)練過(guò)程中不要?jiǎng)討B(tài)增刪requires_grad的凍結(jié)范圍否則每個(gè)step之間的行為會(huì)出現(xiàn)不可復(fù)現(xiàn)的差異。如果你的顯存足夠而DistributedDataParallel在凍結(jié)部分層的條件下依然會(huì)因?yàn)椤澳承﹨?shù)雖然requires_gradFalse但依然侵入同步邏輯”而產(chǎn)生警告可以設(shè)置find_unused_parametersTrue但要注意這會(huì)在每次前向時(shí)多花一點(diǎn)點(diǎn)時(shí)間做探測(cè)。這不是安全問(wèn)題但會(huì)讓你多排查一會(huì)兒。4.4 凍結(jié)后驗(yàn)證參數(shù)真的“凍住”了有時(shí)候你按照教程寫了凍結(jié)代碼模型跑起來(lái)也挺正常卻拿到的結(jié)果不對(duì)。別急著調(diào)整超參數(shù)先花10秒鐘確認(rèn)參數(shù)到底有沒(méi)有在變。一個(gè)非常實(shí)用的檢查手段是記錄訓(xùn)練前后參數(shù)矩陣的差異def check_frozen_param_change(model, before_state, after_state, name_prefix): changed [] for name, param_before in before_state.items(): if name.startswith(name_prefix): param_after after_state[name] if not torch.equal(param_before, param_after): changed.append(name) return changed這能幫你迅速排查到底是哪一層參數(shù)被動(dòng)過(guò)。另一個(gè)快速的檢查手段是訓(xùn)練幾百個(gè)step之后直接查看某個(gè)凍結(jié)層權(quán)重的數(shù)值變化有多小。如果在雙精度下依然完全沒(méi)有變化那說(shuō)明確實(shí)凍住了如果有微小變化說(shuō)明它其實(shí)還在更新得去優(yōu)化器和requires_grad上找原因。5. 常見(jiàn)問(wèn)題速查表與排查指南最后把這些年實(shí)際遇到過(guò)的問(wèn)題整理成一個(gè)速查表方便你直接對(duì)照定位?，F(xiàn)象常見(jiàn)原因排查方法設(shè)置了requires_gradFalse但參數(shù)還在變化優(yōu)化器是之前創(chuàng)建好的仍包含這些參數(shù)引用重新創(chuàng)建優(yōu)化器只傳入requires_gradTrue的參數(shù)模型訓(xùn)練表現(xiàn)正常驗(yàn)證指標(biāo)卻很差BN層還在用當(dāng)前batch統(tǒng)計(jì)量更新running_mean將含BN的模塊切到eval()模式并凍結(jié)其可學(xué)習(xí)參數(shù)load_state_dict報(bào)size mismatchfc等層結(jié)構(gòu)調(diào)整后和預(yù)訓(xùn)練權(quán)重形狀不匹配使用strictFalse并手動(dòng)確認(rèn)能匹配的層數(shù)量多卡訓(xùn)練時(shí)某個(gè)batch的loss不正常各卡凍結(jié)范圍不一致梯度同步出問(wèn)題保證所有進(jìn)程的凍結(jié)邏輯一致避免動(dòng)態(tài)增刪凍結(jié)層明明把某個(gè)層凍住了但顯存沒(méi)有降低凍結(jié)并不會(huì)修改前向傳播路徑激活值仍會(huì)緩存想省顯存需要結(jié)合torch.no_grad()或重寫前向邏輯凍結(jié)后訓(xùn)練不收斂分類頭學(xué)習(xí)率過(guò)大或凍結(jié)范圍過(guò)大導(dǎo)致模型表達(dá)能力不足調(diào)低學(xué)習(xí)率或者部分解凍backbone高層換數(shù)據(jù)集后效果反而更差預(yù)訓(xùn)練特征與目標(biāo)領(lǐng)域差異較大增加解凍層范圍或改用更大學(xué)習(xí)率微調(diào)高層驗(yàn)證時(shí)使用了隨機(jī)初始化參數(shù)忘記加載預(yù)訓(xùn)練權(quán)重就開(kāi)始訓(xùn)練初始化加載后再凍結(jié)并訓(xùn)練這張表看起來(lái)簡(jiǎn)潔但每一個(gè)現(xiàn)象背后我都交過(guò)“學(xué)費(fèi)”。尤其是BN層那條我早期做遷移學(xué)習(xí)時(shí)反反復(fù)復(fù)被它折磨后來(lái)才意識(shí)到requires_gradFalse和BN.eval()是兩件必須同時(shí)做的事。另一個(gè)值得單獨(dú)提一下的建議是在凍結(jié)參數(shù)時(shí)把訓(xùn)練過(guò)程分成熱身階段再進(jìn)入正軌。比如先只訓(xùn)練分類頭10個(gè)epoch再用很小的學(xué)習(xí)率解凍全部參數(shù)微調(diào)5個(gè)epoch很多任務(wù)上效果都會(huì)比從頭全量微調(diào)更穩(wěn)。我自己的經(jīng)驗(yàn)是這種“先冷后熱”的做法通常不額外增加太多訓(xùn)練耗時(shí)但收斂效果和最終指標(biāo)都會(huì)明顯更穩(wěn)。最后再分享一個(gè)小技巧如果你在跑一個(gè)現(xiàn)有項(xiàng)目的代碼別人已經(jīng)把凍結(jié)邏輯寫好了但你并不確定它凍結(jié)了哪些層那就直接在第一個(gè)step之后把全部參數(shù)快照打印出來(lái)就是前面那個(gè)對(duì)比函數(shù)一眼就能看出實(shí)際生效范圍。這個(gè)習(xí)慣幫我避免過(guò)至少十次無(wú)效訓(xùn)練。關(guān)于固定網(wǎng)絡(luò)這件事本質(zhì)上不復(fù)雜但也不是“設(shè)一行requires_gradFalse”就完事了。它和優(yōu)化器、BatchNorm、模型結(jié)構(gòu)、加載邏輯、分布式訓(xùn)練全都牽涉在一起。只要你把層次搞清楚遇到問(wèn)題時(shí)按上面這張表對(duì)號(hào)入座絕大多數(shù)坑都能邁過(guò)去。希望這篇分享對(duì)你接下來(lái)的訓(xùn)練項(xiàng)目有幫助。 PyTorch模型微調(diào)實(shí)戰(zhàn)：凍結(jié)層參數(shù)的原理、策略與踩坑指南
不少朋友第一次接觸模型微調(diào)時(shí)都會(huì)遇到一個(gè)困惑為什么大家都在說(shuō)“凍結(jié)網(wǎng)絡(luò)”、“freeze層參數(shù)”這個(gè)詞聽(tīng)起來(lái)很玄乎其實(shí)說(shuō)白了就是讓某些層的權(quán)重在訓(xùn)練過(guò)程中保持不變只更新剩下的部分。PyTorch里做這件事非常靈活可以用requires_grad False控制梯度也可以在優(yōu)化器里手動(dòng)過(guò)濾參數(shù)還可以配合eval()模式處理BatchNorm這類有狀態(tài)層的坑。這篇文章我想把這套東西徹底講明白。我會(huì)從為什么要凍結(jié)參數(shù)講起然后展示幾種常見(jiàn)的凍結(jié)姿勢(shì)再結(jié)合一個(gè)貼近真實(shí)業(yè)務(wù)的多任務(wù)模型例子把“凍結(jié)哪些層、怎么解凍、學(xué)習(xí)率怎么配”這些實(shí)操細(xì)節(jié)全部過(guò)一遍。你可能是剛?cè)腴TPyTorch的新手也可能已經(jīng)在做遷移學(xué)習(xí)但總被BN層坑得頭皮發(fā)麻這篇文章的思路和踩坑經(jīng)驗(yàn)對(duì)你應(yīng)該都有直接幫助。1. 為什么要凍結(jié)網(wǎng)絡(luò)參數(shù)從省顯存到防過(guò)擬合先說(shuō)結(jié)論凍結(jié)參數(shù)這個(gè)操作本質(zhì)上是在干預(yù)優(yōu)化過(guò)程。神經(jīng)網(wǎng)絡(luò)靠反向傳播更新參數(shù)每一步都需要計(jì)算損失函數(shù)對(duì)各層權(quán)重的梯度。如果把某些層的requires_grad設(shè)為False梯度根本不會(huì)流到那里參數(shù)自然保持不動(dòng)。這一招最常見(jiàn)的應(yīng)用場(chǎng)景有三個(gè)。第一個(gè)場(chǎng)景是遷移學(xué)習(xí)。你拿一個(gè)在大規(guī)模數(shù)據(jù)集上預(yù)訓(xùn)練好的模型當(dāng)特征提取器比如圖像分類里的ResNet、文本分類里的BERT這些模型底層已經(jīng)學(xué)會(huì)了通用的視覺(jué)紋理或句法結(jié)構(gòu)。下游任務(wù)的訓(xùn)練數(shù)據(jù)往往只有幾千張圖、幾萬(wàn)條文本如果從零開(kāi)始訓(xùn)練所有層模型很容易把訓(xùn)練集里的噪聲背下來(lái)。把預(yù)訓(xùn)練部分的參數(shù)凍住只訓(xùn)練最后新增的分類頭相當(dāng)于站在巨人肩膀上做微調(diào)。這個(gè)思路在數(shù)據(jù)量小、任務(wù)與預(yù)訓(xùn)練任務(wù)相似度高的場(chǎng)景下尤其好用。第二個(gè)場(chǎng)景是顯存和算力限制。訓(xùn)練一個(gè)大規(guī)模模型顯存消耗主要來(lái)自中間激活值和優(yōu)化器狀態(tài)。凍結(jié)部分層之后PyTorch在前向傳播時(shí)依然需要算這些層的輸出但反向傳播不再為它們緩存激活值優(yōu)化器也不再為這些層保存動(dòng)量或二階矩估計(jì)。我實(shí)測(cè)過(guò)一個(gè)約3億參數(shù)量的模型凍結(jié)掉底部80%的層之后訓(xùn)練峰值顯存直接降了將近一半單卡跑批大小也能翻倍。第三個(gè)場(chǎng)景容易被人忽略——防止災(zāi)難性遺忘。在做增量學(xué)習(xí)或者多任務(wù)學(xué)習(xí)時(shí)如果希望模型在完成新任務(wù)的同時(shí)不丟掉舊任務(wù)的能力凍結(jié)掉共享底層參數(shù)是性價(jià)比最高的方案之一。我在某個(gè)多語(yǔ)言文本分類項(xiàng)目中一開(kāi)始讓所有層都參與訓(xùn)練結(jié)果發(fā)現(xiàn)模型在新語(yǔ)言上的指標(biāo)上去了舊語(yǔ)言的F1值卻掉得厲害。后來(lái)把共享編碼器的前幾層凍住只訓(xùn)練頂層和任務(wù)專屬的分類分支舊任務(wù)的表現(xiàn)就穩(wěn)定下來(lái)了。當(dāng)然凍結(jié)參數(shù)不是萬(wàn)能的它也有明顯的適用邊界。比如預(yù)訓(xùn)練模型和下游任務(wù)差異巨大時(shí)底層特征完全不更新會(huì)導(dǎo)致模型學(xué)不到任務(wù)相關(guān)的特定表示。再比如你對(duì)某個(gè)任務(wù)的標(biāo)注數(shù)據(jù)足夠多那讓所有層都參與微調(diào)通常能拿到更好的最終精度。所以實(shí)際操作中我一般不會(huì)一上來(lái)就把模型焊死而是先做幾組“全量微調(diào) vs 部分凍結(jié)”的對(duì)照實(shí)驗(yàn)看指標(biāo)差距再?zèng)Q定。2. 核心機(jī)制拆解requires_grad、優(yōu)化器過(guò)濾與訓(xùn)練/評(píng)估模式要把凍結(jié)參數(shù)玩明白先得弄懂PyTorch底層那套梯度傳播邏輯。很多人只知道設(shè)requires_gradFalse就行但對(duì)它背后的機(jī)制、和eval()的搭配、還有優(yōu)化器行為之間的關(guān)系并不清楚結(jié)果總是遇到“明明凍住了參數(shù)怎么還在變”這種怪事。2.1requires_gradFalse的傳播阻斷原理PyTorch的自動(dòng)求導(dǎo)引擎基于計(jì)算圖工作。每當(dāng)你做一次前向傳播框架會(huì)記錄哪些張量參與了運(yùn)算反向傳播時(shí)沿著這個(gè)圖把梯度從損失函數(shù)一路傳回各層的權(quán)重矩陣。requires_gradFalse的作用非常直接這個(gè)張量不會(huì)再出現(xiàn)在計(jì)算圖的葉子節(jié)點(diǎn)列表里反向傳播經(jīng)過(guò)它時(shí)直接跳過(guò)梯度傳不到它身上參數(shù)自然紋絲不動(dòng)。這里有個(gè)細(xì)節(jié)值得說(shuō)透——即使你凍結(jié)了某一層前向傳播仍然會(huì)正常計(jì)算該層的輸出結(jié)果這個(gè)結(jié)果會(huì)被繼續(xù)送入后續(xù)層。也就是說(shuō)網(wǎng)絡(luò)結(jié)構(gòu)沒(méi)有變“前向推理路徑”是完整的只是“參數(shù)的更新路徑”被切斷了。我在給別人講這個(gè)概念時(shí)經(jīng)常打一個(gè)比方凍結(jié)參數(shù)相當(dāng)于把一臺(tái)機(jī)器里某個(gè)調(diào)節(jié)旋鈕擰緊并鎖死機(jī)器還在運(yùn)轉(zhuǎn)產(chǎn)品還是會(huì)出來(lái)只是這個(gè)旋鈕以后不會(huì)再被工人調(diào)整了。寫代碼時(shí)要注意requires_grad必須遞歸設(shè)置到所有子模塊上。比如你用了torchvision.models.resnet50想凍結(jié)整個(gè)特征提取器只寫model.layer1.requires_grad_(False)可不夠因?yàn)閘ayer1里面還有若干個(gè)Bottleneck模塊每個(gè)模塊里面又有多個(gè)卷積層和BN層。正確的做法是遍歷模塊遞歸設(shè)置def freeze_module(module): for param in module.parameters(): param.requires_grad False這個(gè)函數(shù)會(huì)遍歷傳入模塊的所有子模塊并凍結(jié)其參數(shù)比直接操作某個(gè)模塊要穩(wěn)妥得多。2.2 優(yōu)化器里的“隱形坑”參數(shù)過(guò)濾與緩存狀態(tài)很多人在凍結(jié)參數(shù)后忽略了一個(gè)關(guān)鍵步驟——優(yōu)化器需要重新構(gòu)建。假設(shè)你先創(chuàng)建了一個(gè)優(yōu)化器傳入所有參數(shù)然后才把某些層設(shè)成requires_gradFalse那么優(yōu)化器內(nèi)部依然保留著這些參數(shù)的引用并且會(huì)在每個(gè)step里試圖更新它們。雖然由于梯度為0參數(shù)實(shí)際上不會(huì)變化但這會(huì)白白浪費(fèi)計(jì)算和顯存而且沒(méi)有意義的動(dòng)量狀態(tài)累積起來(lái)還可能留下隱患。正確做法是先用requires_gradFalse凍結(jié)參數(shù)再創(chuàng)建優(yōu)化器并且只把那些仍然requires_gradTrue的參數(shù)傳進(jìn)去frozen_params [] trainable_params [] for name, param in model.named_parameters(): if param.requires_grad: trainable_params.append(param) else: frozen_params.append(param) optimizer optim.Adam(trainable_params, lr1e-4)這里我用named_parameters()而不是parameters()目的是拿到每個(gè)參數(shù)的名字方便后續(xù)打印、篩選或者按層名稱做精細(xì)化的學(xué)習(xí)率分組。而且把這個(gè)參數(shù)列表打印出來(lái)你會(huì)很清楚當(dāng)前到底哪些層在訓(xùn)練、哪些層是凍結(jié)的。另一個(gè)容易踩的坑是優(yōu)化器緩存。某些優(yōu)化器比如SGD的動(dòng)量、Adam的一階二階動(dòng)量是根據(jù)參數(shù)狀態(tài)來(lái)維護(hù)緩存的。你雖然已經(jīng)設(shè)了requires_gradFalse但舊優(yōu)化器里關(guān)于這些參數(shù)的緩存可能依然存在。更穩(wěn)妥的做法是只要你調(diào)整了凍結(jié)范圍就直接重新創(chuàng)建一個(gè)優(yōu)化器實(shí)例不要復(fù)用舊的。這不用寫多少代碼但能有效避免很多詭異的問(wèn)題。2.3train()與eval()模式不只是Dropout的開(kāi)關(guān)初學(xué)者很容易把參數(shù)凍結(jié)和model.eval()混為一談。它倆的確有交集但完全是兩碼事。model.eval()會(huì)把模型切到評(píng)估模式影響的是Dropout、BatchNorm這類在訓(xùn)練和推理時(shí)行為不同的層但它不會(huì)把任何參數(shù)設(shè)成requires_gradFalse。也就是說(shuō)即使你調(diào)用了model.eval()只要優(yōu)化器里還有這些參數(shù)它們照樣會(huì)被更新。反過(guò)來(lái)凍結(jié)參數(shù)也不意味著你要把整個(gè)模型切到eval()模式。這里有個(gè)典型誤區(qū)很多人凍結(jié)了特征提取層之后習(xí)慣性地調(diào)用了model.eval()結(jié)果BN層在推理時(shí)用的是全局統(tǒng)計(jì)量而不是當(dāng)前batch的統(tǒng)計(jì)量這反而影響了訓(xùn)練效果。更常見(jiàn)的做法是只對(duì)含BatchNorm的部分單獨(dú)處理或者干脆將BN層也放在訓(xùn)練模式下但凍結(jié)其中的可學(xué)習(xí)參數(shù)。在動(dòng)手之前請(qǐng)務(wù)必弄清楚這幾件事的邊界requires_gradFalse控制的是梯度計(jì)算和參數(shù)更新optimizer決定哪些參數(shù)參與更新model.train()/model.eval()控制的是層的行為模式。參數(shù)凍結(jié)時(shí)我一般會(huì)把最終要訓(xùn)練的那部分模塊保持train()模式但會(huì)根據(jù)BN層是否會(huì)被凍結(jié)來(lái)單獨(dú)處理。3. 實(shí)操?gòu)囊粋€(gè)多任務(wù)模型看不同凍結(jié)策略的代碼落地理論說(shuō)得再多不如看一個(gè)完整的實(shí)戰(zhàn)例子。我在很多業(yè)務(wù)場(chǎng)景中會(huì)遇到這樣的情況模型需要同時(shí)完成一個(gè)主任務(wù)和一個(gè)輔助任務(wù)但輔助任務(wù)的標(biāo)注數(shù)據(jù)非常有限如果讓所有參數(shù)都參與優(yōu)化很容易在輔助任務(wù)上過(guò)擬合。下面我用一個(gè)簡(jiǎn)化版本的多任務(wù)分類模型來(lái)演示不同的凍結(jié)策略。3.1 搭建一個(gè)共享底層、多分支輸出的模型假設(shè)我們有一個(gè)眾包標(biāo)注平臺(tái)需要同時(shí)預(yù)測(cè)圖片的“主類別”和“質(zhì)量分?jǐn)?shù)”。主類別的標(biāo)注數(shù)據(jù)很充裕質(zhì)量分?jǐn)?shù)的標(biāo)注卻比較少。一個(gè)合理的模型結(jié)構(gòu)是共享一個(gè)ResNet風(fēng)格的編碼器然后分出兩個(gè)分類頭。這個(gè)模擬項(xiàng)目X的目標(biāo)就是驗(yàn)證“凍結(jié)共享編碼器、只訓(xùn)練兩個(gè)分類頭”能否在輔助任務(wù)上取得更好的泛化效果。import torch import torch.nn as nn from torchvision import models class MultiTaskModel(nn.Module): def __init__(self): super().__init__() self.backbone models.resnet18(pretrainedTrue) self.backbone.fc nn.Identity() # 去掉原來(lái)的全連接層 self.head_main nn.Sequential( nn.Linear(512, 128), nn.ReLU(inplaceTrue), nn.Linear(128, 10) ) self.head_aux nn.Sequential( nn.Linear(512, 64), nn.ReLU(inplaceTrue), nn.Linear(64, 1) ) def forward(self, x): feat self.backbone(x) out_main self.head_main(feat) out_aux self.head_aux(feat) return out_main, out_aux這個(gè)模型很典型共享編碼器是預(yù)訓(xùn)練好的ResNet18兩個(gè)任務(wù)各有一個(gè)小分類頭。現(xiàn)在我們要把策略落在代碼上。3.2 策略一只訓(xùn)練分類頭凍結(jié)全部backbone參數(shù)這個(gè)策略最簡(jiǎn)單也是我早期最常用的方法。整個(gè)預(yù)訓(xùn)練編碼器一筆不動(dòng)只訓(xùn)練兩個(gè)新增的全連接分支。這樣做的好處是訓(xùn)練速度極快baseline跑起來(lái)特別快幾分鐘就能看到路線是否可行。def freeze_encoder_and_train_heads(model): # 凍結(jié)backbone for name, param in model.backbone.named_parameters(): param.requires_grad False # 只把head里的參數(shù)傳給優(yōu)化器 trainable_params [] for name, param in model.named_parameters(): if param.requires_grad: trainable_params.append(param) optimizer torch.optim.Adam(trainable_params, lr1e-3) return optimizer這里有個(gè)很容易被忽略的細(xì)節(jié)我們用的是model.named_parameters()但param.requires_grad是直接讀的參數(shù)屬性。如果某些層之前已經(jīng)被設(shè)置為不可訓(xùn)練這里會(huì)自動(dòng)跳過(guò)它們無(wú)需額外過(guò)濾。訓(xùn)練時(shí)你可以順手打印當(dāng)前正在訓(xùn)練的層名確認(rèn)自己的意圖真正落到了代碼層面for name, param in model.named_parameters(): if param.requires_grad: print(訓(xùn)練中的參數(shù):, name)3.3 策略二凍結(jié)backbone的前幾層解凍后面高層并設(shè)置差異化學(xué)習(xí)率凍結(jié)全部backbone最省事但上限也不高。如果任務(wù)和預(yù)訓(xùn)練任務(wù)差距比較大底層特征完全不動(dòng)的話最終精度往往不夠理想。實(shí)踐中我會(huì)用一個(gè)更精細(xì)的策略底層的淺層特征比如邊緣、顏色、紋理這些通用結(jié)構(gòu)完全凍結(jié)高層的語(yǔ)義特征比如跟具體物體部件相關(guān)的特征解凍微調(diào)兩個(gè)分類頭用更高的學(xué)習(xí)率backbone的高層用較低學(xué)習(xí)率。def partial_freeze_with_lr(model): # 凍結(jié)layer1和layer2解凍layer3和layer4 for name, param in model.backbone.named_parameters(): if name.startswith(layer1) or name.startswith(layer2): param.requires_grad False else: param.requires_grad True optimizer torch.optim.Adam([ # 分類頭大學(xué)習(xí)率 {params: model.head_main.parameters(), lr: 1e-3}, {params: model.head_aux.parameters(), lr: 1e-3}, # backbone高層小學(xué)習(xí)率 {params: filter(lambda p: p.requires_grad, model.backbone.layer3.parameters()), lr: 1e-5}, {params: filter(lambda p: p.requires_grad, model.backbone.layer4.parameters()), lr: 1e-5}, ]) return optimizer差異化學(xué)習(xí)率是這策略的靈魂。底層特征通用性強(qiáng)、更新風(fēng)險(xiǎn)高就完全鎖死高層特征與具體任務(wù)相關(guān)性強(qiáng)可以適當(dāng)更新但預(yù)訓(xùn)練的知識(shí)也別丟得太多所以用一個(gè)很小的學(xué)習(xí)率。分類頭是隨機(jī)初始化的要從頭學(xué)所以給它最大的學(xué)習(xí)率。這樣一組參數(shù)在分類準(zhǔn)確率和訓(xùn)練穩(wěn)定性上都比較平衡。關(guān)于解凍的邊界一個(gè)穩(wěn)妥的方法是觀察每層特征圖的激活變化。你可以拿一個(gè)固定batch的數(shù)據(jù)在前向傳播時(shí)打印每層輸出的均值方差對(duì)比凍結(jié)前后哪些層的輸出變化最大。如果發(fā)現(xiàn)layer2解凍后輸出分布已經(jīng)亂得不成樣子說(shuō)明學(xué)習(xí)率太高了或者干脆把它凍回去更穩(wěn)。3.4 策略三在凍結(jié)狀態(tài)下做“冷啟動(dòng)”再用“解凍全網(wǎng)微調(diào)”做熱階段有些任務(wù)雖然數(shù)據(jù)量一般但遷移學(xué)習(xí)的起點(diǎn)非常關(guān)鍵。我慣用的一個(gè)方法是兩階段訓(xùn)練第一階段讓backbone全面凍結(jié)只訓(xùn)練分類頭目標(biāo)是讓分類頭在預(yù)訓(xùn)練特征上找到一個(gè)大致的解第二階段解凍全部參數(shù)用一個(gè)非常小的學(xué)習(xí)率做全網(wǎng)微調(diào)。這比直接從隨機(jī)初始化訓(xùn)練整個(gè)網(wǎng)絡(luò)要穩(wěn)定不少也能明顯提升最終精度。# 階段一只訓(xùn)練head def stage_one(model): for p in model.backbone.parameters(): p.requires_grad False return torch.optim.Adam(model.head_main.parameters(), lr1e-3) # 階段二全部解凍 低學(xué)習(xí)率 def stage_two(model): for p in model.backbone.parameters(): p.requires_grad True optimizer torch.optim.Adam(model.parameters(), lr1e-5) return optimizer這里的關(guān)鍵點(diǎn)是階段一和階段二之間除了學(xué)習(xí)率變化最好把階段一的優(yōu)化器丟棄掉重新建一個(gè)新的。原因前面提過(guò)優(yōu)化器緩存里還保存著階段一的狀態(tài)突然換到全網(wǎng)微調(diào)時(shí)那些狀態(tài)可能會(huì)干擾更新方向重新創(chuàng)建更干凈。4. BN層、參數(shù)加載與多卡訓(xùn)練的連環(huán)坑凍結(jié)參數(shù)看起來(lái)是個(gè)小操作但在真實(shí)項(xiàng)目里總會(huì)碰到各種“看著沒(méi)事、跑起來(lái)就掛”的情況。這里面的坑我基本都踩過(guò)一遍一個(gè)個(gè)說(shuō)。4.1 BatchNorm層只凍結(jié)可學(xué)習(xí)參數(shù)遠(yuǎn)遠(yuǎn)不夠這是最容易踩、也最隱蔽的坑。BatchNorm層雖然是通過(guò)weight和bias學(xué)習(xí)縮放和平移但它的行為還依賴兩個(gè)“狀態(tài)量”——running_mean和running_var。這兩個(gè)量不是靠反向傳播更新的而是在每次前向傳播時(shí)用一個(gè)動(dòng)量系數(shù)在線更新的。所以即使你設(shè)了requires_gradFalse只要model.train()模式下這個(gè)BN層還在跑前向計(jì)算running_mean和running_var就會(huì)不斷被當(dāng)前batch的統(tǒng)計(jì)量更新。這在訓(xùn)練和推理時(shí)會(huì)出現(xiàn)嚴(yán)重的分布漂移問(wèn)題。如何正確凍結(jié)BatchNorm層我給你一個(gè)標(biāo)準(zhǔn)做法def freeze_bn(mod): for m in mod.modules(): if isinstance(m, nn.BatchNorm2d): m.eval() # 讓BN層使用全局統(tǒng)計(jì)量而不是更新running_mean for p in m.parameters(): p.requires_grad False這樣做的本質(zhì)是讓BN層切到評(píng)測(cè)模式前向傳播時(shí)直接使用running_mean和running_var不再用本batch的統(tǒng)計(jì)量更新它們。這樣才能做到真正的“凍結(jié)”。很多人只設(shè)了requires_gradFalse卻忘了讓BN層切模式結(jié)果模型在訓(xùn)練時(shí)表現(xiàn)正常、在驗(yàn)證時(shí)指標(biāo)掉得厲害八成就是這里出了問(wèn)題。那到底要凍結(jié)BN嗎需要分場(chǎng)景討論。如果凍結(jié)的是整個(gè)backboneBN層務(wù)必按上面這種方式處理因?yàn)樗鼉?nèi)部的統(tǒng)計(jì)數(shù)據(jù)如果持續(xù)變化會(huì)干擾凍結(jié)層輸出的特征一致性如果模型很小、任務(wù)分布和預(yù)訓(xùn)練分布差別很大BN也可以保持默認(rèn)的跟蹤更新只是這時(shí)它沒(méi)有被真正凍結(jié)你需要心里有數(shù)。4.2 加載預(yù)訓(xùn)練權(quán)重時(shí)嚴(yán)格對(duì)齊參數(shù)名和形狀用別人訓(xùn)練好的權(quán)重來(lái)微調(diào)如果你的模型結(jié)構(gòu)是從torchvision.models.resnet18(pretrainedTrue)直接加載的那名字天然對(duì)得上。但如果你自己改了結(jié)構(gòu)比如把fc換成一個(gè)雙分支頭加載權(quán)重時(shí)就很可能報(bào)size mismatch的錯(cuò)。這類報(bào)錯(cuò)原因很直接預(yù)訓(xùn)練權(quán)重的某個(gè)層形狀和你的模型對(duì)不上。比如原始的ResNet18最后一個(gè)全連接層是(512, 1000)形狀你把它換成了(512, 10)那加載時(shí)就肯定會(huì)沖突。解決辦法是strictFalse只加載能對(duì)上的部分checkpoint torch.load(pretrained_weights.pth) model.load_state_dict(checkpoint, strictFalse)但這里要小心strictFalse會(huì)把無(wú)法匹配的層全部忽略包括你本來(lái)想加載的層。要確保想加載的特征提取層能被正確覆蓋我建議加載前先打印一下檢查一下model.load_state_dict(checkpoint, strictFalse) for name, param in model.named_parameters(): print(name, param.shape, param.requires_grad)另外如果你是自己動(dòng)手做的模型加載之前還應(yīng)該確認(rèn)state_dict里的鍵名是否跟模型里的一致。比如某層以前叫backbone.conv1.weight后來(lái)你重構(gòu)成了encoder.conv1.weight鍵名就對(duì)不上了得手工做一下映射。最省事的做法是在模型保存和加載時(shí)都使用同一套命名規(guī)則別在中間改名。4.3 多卡訓(xùn)練時(shí)凍結(jié)邏輯與All-Reduce的交互使用DataParallel或DistributedDataParallel時(shí)凍結(jié)參數(shù)還有一個(gè)需要提前想清楚的地方?？蚣茉诿看畏聪騻鞑r(shí)會(huì)對(duì)梯度做跨卡同步All-Reduce。如果你只在一張卡上凍結(jié)參數(shù)而在其他卡上沒(méi)有保持一致梯度同步就會(huì)出問(wèn)題最終導(dǎo)致更新行為完全不可控。而且要特別注意在DistributedDataParallel模式下歸一化層和同步層本來(lái)就是全局同步的。如果通過(guò)某些方式把BN層也帶上了同步邏輯你凍結(jié)時(shí)也要考慮sync BN的行為。我的經(jīng)驗(yàn)是在做多卡訓(xùn)練前先把凍結(jié)參數(shù)的全部邏輯固定下來(lái)讓所有進(jìn)程在完全一致的參數(shù)狀態(tài)下進(jìn)入訓(xùn)練在訓(xùn)練過(guò)程中不要?jiǎng)討B(tài)增刪requires_grad的凍結(jié)范圍否則每個(gè)step之間的行為會(huì)出現(xiàn)不可復(fù)現(xiàn)的差異。如果你的顯存足夠而DistributedDataParallel在凍結(jié)部分層的條件下依然會(huì)因?yàn)椤澳承﹨?shù)雖然requires_gradFalse但依然侵入同步邏輯”而產(chǎn)生警告可以設(shè)置find_unused_parametersTrue但要注意這會(huì)在每次前向時(shí)多花一點(diǎn)點(diǎn)時(shí)間做探測(cè)。這不是安全問(wèn)題但會(huì)讓你多排查一會(huì)兒。4.4 凍結(jié)后驗(yàn)證參數(shù)真的“凍住”了有時(shí)候你按照教程寫了凍結(jié)代碼模型跑起來(lái)也挺正常卻拿到的結(jié)果不對(duì)。別急著調(diào)整超參數(shù)先花10秒鐘確認(rèn)參數(shù)到底有沒(méi)有在變。一個(gè)非常實(shí)用的檢查手段是記錄訓(xùn)練前后參數(shù)矩陣的差異def check_frozen_param_change(model, before_state, after_state, name_prefix): changed [] for name, param_before in before_state.items(): if name.startswith(name_prefix): param_after after_state[name] if not torch.equal(param_before, param_after): changed.append(name) return changed這能幫你迅速排查到底是哪一層參數(shù)被動(dòng)過(guò)。另一個(gè)快速的檢查手段是訓(xùn)練幾百個(gè)step之后直接查看某個(gè)凍結(jié)層權(quán)重的數(shù)值變化有多小。如果在雙精度下依然完全沒(méi)有變化那說(shuō)明確實(shí)凍住了如果有微小變化說(shuō)明它其實(shí)還在更新得去優(yōu)化器和requires_grad上找原因。5. 常見(jiàn)問(wèn)題速查表與排查指南最后把這些年實(shí)際遇到過(guò)的問(wèn)題整理成一個(gè)速查表方便你直接對(duì)照定位?，F(xiàn)象常見(jiàn)原因排查方法設(shè)置了requires_gradFalse但參數(shù)還在變化優(yōu)化器是之前創(chuàng)建好的仍包含這些參數(shù)引用重新創(chuàng)建優(yōu)化器只傳入requires_gradTrue的參數(shù)模型訓(xùn)練表現(xiàn)正常驗(yàn)證指標(biāo)卻很差BN層還在用當(dāng)前batch統(tǒng)計(jì)量更新running_mean將含BN的模塊切到eval()模式并凍結(jié)其可學(xué)習(xí)參數(shù)load_state_dict報(bào)size mismatchfc等層結(jié)構(gòu)調(diào)整后和預(yù)訓(xùn)練權(quán)重形狀不匹配使用strictFalse并手動(dòng)確認(rèn)能匹配的層數(shù)量多卡訓(xùn)練時(shí)某個(gè)batch的loss不正常各卡凍結(jié)范圍不一致梯度同步出問(wèn)題保證所有進(jìn)程的凍結(jié)邏輯一致避免動(dòng)態(tài)增刪凍結(jié)層明明把某個(gè)層凍住了但顯存沒(méi)有降低凍結(jié)并不會(huì)修改前向傳播路徑激活值仍會(huì)緩存想省顯存需要結(jié)合torch.no_grad()或重寫前向邏輯凍結(jié)后訓(xùn)練不收斂分類頭學(xué)習(xí)率過(guò)大或凍結(jié)范圍過(guò)大導(dǎo)致模型表達(dá)能力不足調(diào)低學(xué)習(xí)率或者部分解凍backbone高層換數(shù)據(jù)集后效果反而更差預(yù)訓(xùn)練特征與目標(biāo)領(lǐng)域差異較大增加解凍層范圍或改用更大學(xué)習(xí)率微調(diào)高層驗(yàn)證時(shí)使用了隨機(jī)初始化參數(shù)忘記加載預(yù)訓(xùn)練權(quán)重就開(kāi)始訓(xùn)練初始化加載后再凍結(jié)并訓(xùn)練這張表看起來(lái)簡(jiǎn)潔但每一個(gè)現(xiàn)象背后我都交過(guò)“學(xué)費(fèi)”。尤其是BN層那條我早期做遷移學(xué)習(xí)時(shí)反反復(fù)復(fù)被它折磨后來(lái)才意識(shí)到requires_gradFalse和BN.eval()是兩件必須同時(shí)做的事。另一個(gè)值得單獨(dú)提一下的建議是在凍結(jié)參數(shù)時(shí)把訓(xùn)練過(guò)程分成熱身階段再進(jìn)入正軌。比如先只訓(xùn)練分類頭10個(gè)epoch再用很小的學(xué)習(xí)率解凍全部參數(shù)微調(diào)5個(gè)epoch很多任務(wù)上效果都會(huì)比從頭全量微調(diào)更穩(wěn)。我自己的經(jīng)驗(yàn)是這種“先冷后熱”的做法通常不額外增加太多訓(xùn)練耗時(shí)但收斂效果和最終指標(biāo)都會(huì)明顯更穩(wěn)。最后再分享一個(gè)小技巧如果你在跑一個(gè)現(xiàn)有項(xiàng)目的代碼別人已經(jīng)把凍結(jié)邏輯寫好了但你并不確定它凍結(jié)了哪些層那就直接在第一個(gè)step之后把全部參數(shù)快照打印出來(lái)就是前面那個(gè)對(duì)比函數(shù)一眼就能看出實(shí)際生效范圍。這個(gè)習(xí)慣幫我避免過(guò)至少十次無(wú)效訓(xùn)練。關(guān)于固定網(wǎng)絡(luò)這件事本質(zhì)上不復(fù)雜但也不是“設(shè)一行requires_gradFalse”就完事了。它和優(yōu)化器、BatchNorm、模型結(jié)構(gòu)、加載邏輯、分布式訓(xùn)練全都牽涉在一起。只要你把層次搞清楚遇到問(wèn)題時(shí)按上面這張表對(duì)號(hào)入座絕大多數(shù)坑都能邁過(guò)去。希望這篇分享對(duì)你接下來(lái)的訓(xùn)練項(xiàng)目有幫助。