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