優(yōu)化器是 PyTorch 訓(xùn)練流程里最容易被“抄作業(yè)”抄錯的部分。很多人跑模型的時候直接照搬別人代碼里的optim.Adam(model.parameters(), lr3e-4)換了個任務(wù)就發(fā)現(xiàn) loss 不降、梯度爆炸、收斂慢半拍問題往往就出在優(yōu)化器選型和參數(shù)配置上。這幾個月我在各種實驗里反復(fù)折騰torch.optim從最基礎(chǔ)的 SGD 到 AdamW、從學(xué)習(xí)率調(diào)度到梯度裁剪踩了不少坑也把相關(guān)論文翻了底朝天。這篇就從一個實戰(zhàn)者的角度把 PyTorch 優(yōu)化器的底細、論文出處、參數(shù)含義和實操踩坑經(jīng)驗一次性講清楚。無論你是剛把 PyTorch 裝好、準備跑第一個模型的入門選手還是已經(jīng)在調(diào)參路上掙扎了一段時間的煉丹學(xué)徒這篇文章應(yīng)該都能幫你少走很多彎路。1. torch.optim 到底是什么先搞清楚優(yōu)化器在訓(xùn)練中的位置1.1 優(yōu)化器解決的是哪一步問題很多人一上來就學(xué)各種網(wǎng)絡(luò)的代碼反而忽略了最基礎(chǔ)的問題優(yōu)化器在整個訓(xùn)練流程里到底干了什么。一句話概括優(yōu)化器負責(zé)回答“參數(shù)往哪個方向走、走多遠”這個問題。神經(jīng)網(wǎng)絡(luò)訓(xùn)練的本質(zhì)是讓損失函數(shù)的值盡可能小而損失函數(shù)是模型參數(shù)的函數(shù)。我們通過反向傳播算出每個參數(shù)對損失的梯度梯度告訴你的是“在這個點上往上走還是往下走才能讓損失變小”。優(yōu)化器拿到這個梯度之后按照某種策略更新模型參數(shù)讓損失一步步下降。這個行為特別像下山損失函數(shù)是連綿起伏的山地模型參數(shù)是你的坐標梯度是腳下山坡的傾斜方向。你要做的就是一步步往低處走。但具體每一步邁多大、怎么避開坑、怎么從局部小坑里爬出來這就是優(yōu)化器策略的學(xué)問。PyTorch 把這一整套邏輯封裝在torch.optim里。你不需要手寫參數(shù)更新公式只需要告訴它“你要優(yōu)化哪些參數(shù)、用什么算法、學(xué)習(xí)率設(shè)多少”。剩下的事情——梯度更新、動量累積、自適應(yīng)學(xué)習(xí)率計算——它全部幫你做了。1.2 torch.optim 的設(shè)計哲學(xué)把優(yōu)化策略做成可插拔模塊我最早用 PyTorch 的時候有一個感受特別深換優(yōu)化器太方便了。從 SGD 切到 Adam只需要改一行代碼訓(xùn)練邏輯、反向傳播、模型結(jié)構(gòu)完全不用動。這個便利性不是巧合而是 PyTorch 把優(yōu)化器規(guī)范化成了統(tǒng)一的接口。torch.optim的核心抽象是Optimizer基類。任何一個優(yōu)化器做的事情都可以拆成兩步根據(jù)當(dāng)前梯度計算出參數(shù)更新量然后把更新量作用到參數(shù)上。PyTorch 把這個過程封裝成optimizer.step()訓(xùn)練循環(huán)里每輪迭代調(diào)用一次即可。這個設(shè)計帶來的好處非常實際實驗管理變得極其輕松。我在做對比實驗的時候經(jīng)常在同一份代碼里用argparse傳一個--optim參數(shù)直接切換 SGD、Adam、AdamW、RMSprop完全不需要改動模型代碼。對于需要大規(guī)模調(diào)參的研究項目或工業(yè)項目來說這種可替換性節(jié)省的時間非?？捎^。1.3 一個優(yōu)化器只有三樣?xùn)|西狀態(tài)、參數(shù)、更新規(guī)則別看torch.optim里面優(yōu)化器種類很多拆開來看每個優(yōu)化器其實就包含三個部分。第一是參數(shù)組。優(yōu)化器存的不是模型本身而是模型參數(shù)的引用。你可以選擇優(yōu)化全部參數(shù)也可以只優(yōu)化部分參數(shù)比如凍結(jié)backbone只微調(diào)分類頭。PyTorch 用param_groups這個概念來管理每個參數(shù)組可以有不同的學(xué)習(xí)率、動量等超參數(shù)。第二是狀態(tài)。像 SGD 的動量項、Adam 的一階矩和二階矩估計這些東西是跨 step 累積的中間量PyTorch 把它們存在優(yōu)化器的state字典里。這也是為什么你保存模型權(quán)重之后還要單獨保存優(yōu)化器狀態(tài)才能精確恢復(fù)訓(xùn)練現(xiàn)場。第三是更新規(guī)則。不同的優(yōu)化器本質(zhì)上就是梯度更新公式不同。SGD 直接用梯度乘學(xué)習(xí)率Adam 在梯度基礎(chǔ)上做了指數(shù)滑動平均和自適應(yīng)縮放。這部分就是各個優(yōu)化器源碼里step()方法的核心邏輯。理解這三樣?xùn)|西你再看優(yōu)化器源碼就不會覺得頭暈。我之前為了搞清楚 Adam 的state里到底存了什么直接打印過optimizer.state_dict()看到里面每個參數(shù)的exp_avg和exp_avg_sq就全明白了。2. 主流優(yōu)化器逐一拆解從 SGD 到 AdamW2.1 SGD最樸素的基線但永遠不要小看它隨機梯度下降Stochastic Gradient DescentSGD是深度學(xué)習(xí)的基石。它的更新規(guī)則極其簡單參數(shù)減去學(xué)習(xí)率乘以梯度。用 PyTorch 實現(xiàn)就是一行optimizer torch.optim.SGD(model.parameters(), lr0.01)樸素 SGD 的優(yōu)點是內(nèi)存占用極低、計算開銷極小、每步更新只依賴當(dāng)前梯度。但它有一個致命弱點收斂慢而且容易在 loss 地形中“震蕩”。這個震蕩怎么理解想象你在山谷里往下走如果這個山谷的橫截面是個狹長的陡溝那么下山的最優(yōu)路徑應(yīng)該順著溝的長軸方向走。但樸素 SGD 只看當(dāng)前點的梯度方向這個方向在溝壁兩側(cè)來回擺動導(dǎo)致你不斷橫跳前進速度極慢。這就是病態(tài)曲率問題。所以實際工程里大家用的 SGD 幾乎都是帶動量Momentum的 SGD。動量項相當(dāng)于給更新過程加了“慣性”讓參數(shù)更新方向不只由當(dāng)前梯度決定還參考歷史梯度的累計方向這樣在溝槽地形里就能有效抑制震蕩、加速收斂。optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9)momentum0.9是經(jīng)驗上最常用的默認值。它表示新的更新方向有 90% 的比例來自歷史動量、10% 來自當(dāng)前梯度。這個參數(shù)調(diào)得越高更新越平滑但可能反應(yīng)越遲鈍。PyTorch 里還支持nesterovTrue開啟 Nesterov 動量這是動量的改進版在計算梯度時先“往前看一步”收斂速度通常更快一些。2.2 Adam自適應(yīng)學(xué)習(xí)率的集大成者如果說 SGD 是手動擋那 Adam 就是自動擋。Adam 全稱 Adaptive Moment Estimation它在 2014 年由 Kingma 和 Ba 提出可以說是深度學(xué)習(xí)訓(xùn)練里使用頻率最高的優(yōu)化器。Adam 核心的創(chuàng)新在于為每個參數(shù)單獨維護一個自適應(yīng)學(xué)習(xí)率。具體來說它維護了兩個量一階矩估計梯度的均值相當(dāng)于動量和二階矩估計梯度平方的均值反映了梯度的振幅。二階矩大的參數(shù)說明這個方向梯度波動劇烈更新步子就小一點二階矩小的參數(shù)更新步子可以大一點。PyTorch 里最常用的寫法是optimizer torch.optim.Adam(model.parameters(), lr1e-3, betas(0.9, 0.999), eps1e-8)這套默認參數(shù)是原論文里給出的經(jīng)受住了大量實驗驗證。lr1e-3是我見過的絕大多數(shù)任務(wù)里 Adam 收斂效果不錯的起點。betas是兩個滑動平均的衰減系數(shù)0.9對應(yīng)動量0.999對應(yīng)梯度平方的滑動平均。eps是一個極小的數(shù)用來防止分母除零。Adam 最大的優(yōu)點是不太需要精細調(diào)學(xué)習(xí)率。同樣是lr1e-3在 ResNet、Transformer、GAN 等不同架構(gòu)上通常都能得到一個還能接受的訓(xùn)練結(jié)果。這對初學(xué)者極其友好也解釋了為什么網(wǎng)上大量代碼里默認就是Adam(lr1e-3)。但 Adam 也有自己的毛病。其中最常見的是“極端不收斂”或后期泛化能力差的問題。一些研究表明Adam 使用指數(shù)滑動平均來估計梯度統(tǒng)計量在訓(xùn)練后期這些估計值可能存在偏差導(dǎo)致模型的泛化性能不如精細調(diào)過的 SGD帶動量。這在圖像分類等任務(wù)上表現(xiàn)特別明顯SGD 最終的測試精度往往比 Adam 高一點。2.3 AdamW權(quán)重衰減的正確打開方式AdamW 是我這幾年來最常使用的優(yōu)化器尤其是做 Transformer、BERT 這類模型的時候。它的全稱是 Adam with Decoupled Weight Decay由 Ilya Loshchilov 和 Frank Hutter 在 2019 年提出。要理解 AdamW先要理解什么是權(quán)重衰減。L2 正則化是給損失函數(shù)加上所有權(quán)重的平方和讓權(quán)重不要長得太大。傳統(tǒng)的 Adam 在實現(xiàn) L2 正則化時是把權(quán)重衰減項直接加到梯度里然后再做自適應(yīng)學(xué)習(xí)率計算。但這樣有一個被證實的缺陷自適應(yīng)學(xué)習(xí)率會“扭曲”權(quán)重衰減的效果使得不同權(quán)重被衰減的程度不一致從而削弱正則化的效果。AdamW 的做法是把權(quán)重衰減從梯度計算中解耦出來。在 Adam 計算出更新量之后直接對參數(shù)做一次額外的小幅縮減這個縮減不受自適應(yīng)學(xué)習(xí)率影響。聽起來只是一個小改動但實際效果非常明顯。我在訓(xùn)練一些大規(guī)模模型時從 Adam 換成 AdamW驗證集表現(xiàn)經(jīng)常有可見提升而且更容易訓(xùn)練出泛化性能更好的模型。用法同樣很簡單optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay0.01)weight_decay0.01是我在 HuggingFace 生態(tài)里看到的最常見配置。如果是 fine-tune 預(yù)訓(xùn)練模型很多人也會用weight_decay0.05或0.1。這個值不是越大越好太大會讓模型欠擬合太小則起不到正則化作用。需要根據(jù)任務(wù)規(guī)模做實驗。2.4 RMSprop、Adagrad、Adadelta各有各的看家本領(lǐng)除了 SGD 和 AdamPyTorch 還提供了好幾個優(yōu)化器雖然日常使用頻率不高但在特定場景里依然有一席之地。Adagrad 是最早的自適應(yīng)學(xué)習(xí)率方法之一。它的特點是對低頻出現(xiàn)的參數(shù)給予更大的更新步長很適合處理稀疏特征比如推薦系統(tǒng)里的用戶/物品 ID 嵌入。但它的缺點是累積梯度平方會越來越大導(dǎo)致學(xué)習(xí)率被壓得過小最終訓(xùn)練提前停滯。RMSprop 是 Hinton 在課程中提出的一種優(yōu)化方法它通過指數(shù)滑動平均來估計梯度平方解決了 Adagrad 學(xué)習(xí)率單調(diào)遞減的問題。在強化學(xué)習(xí)里RMSprop 曾經(jīng)是標配優(yōu)化器之一。PyTorch 中的實現(xiàn)提供alpha參數(shù)控制滑動平均的衰減默認0.99。Adadelta 可以看作是 RMSprop 的一個變體它不需要手動設(shè)置學(xué)習(xí)率更新量會自適應(yīng)地縮放。不過實際使用中它的通用性不如 Adam我很少在主流框架里看到它作為首選。這幾個優(yōu)化器的 PyTorch 用法都是同一套optimizer torch.optim.Adagrad(model.parameters(), lr0.01) optimizer torch.optim.RMSprop(model.parameters(), lr0.01, alpha0.99) optimizer torch.optim.Adadelta(model.parameters(), rho0.9)2.5 優(yōu)化器對比一張表看透區(qū)別我一直覺得選優(yōu)化器之前先看一張對比表比自己瞎試高效得多。這里把我常用的幾個優(yōu)化器的關(guān)鍵特性整理出來優(yōu)化器核心機制適用場景默認學(xué)習(xí)率主要缺點SGD直接沿梯度方向更新CNN圖像分類、大規(guī)模數(shù)據(jù)0.01~0.1收斂慢、需精細調(diào)參SGDMomentum梯度歷史動量大部分CNN任務(wù)0.01~0.1仍依賴學(xué)習(xí)率調(diào)度Adam一階矩二階矩自適應(yīng)大多數(shù)場景通用3e-4~1e-3泛化性能有時不如SGDAdamWAdam解耦權(quán)重衰減Transformer、預(yù)訓(xùn)練模型1e-4~1e-3需要設(shè)置weight_decayRMSprop梯度平方滑動平均強化學(xué)習(xí)、RNN1e-3通用性不如AdamAdagrad梯度平方累積稀疏特征場景1e-2學(xué)習(xí)率衰減過快我的經(jīng)驗是如果做圖像分類這種比較“正統(tǒng)”的監(jiān)督學(xué)習(xí)我會優(yōu)先考慮帶動量的 SGD配合余弦學(xué)習(xí)率衰減最終精度通常不錯如果做 NLP、Transformer、GAN 這類對調(diào)參敏感的任務(wù)直接用 AdamW省心且效果好快速驗證某個想法能不能跑通就上 Adam拿默認參數(shù)先跑通再說。3. 優(yōu)化器參數(shù)精細解讀別再亂調(diào)了3.1 learning rate整個訓(xùn)練里最重要的超參數(shù)學(xué)習(xí)率是優(yōu)化器里最核心、也最需要謹慎調(diào)整的參數(shù)。它直接決定每一步更新跨多大步。學(xué)習(xí)率設(shè)得太大參數(shù)會在最優(yōu)解附近來回震蕩甚至發(fā)散設(shè)得太小訓(xùn)練過程極其緩慢且容易陷入局部最優(yōu)出不來。一個我常用的經(jīng)驗判斷訓(xùn)練開始時如果 loss 在下降但非常慢可以試著增大學(xué)習(xí)率如果 loss 出現(xiàn)劇烈波動甚至變成 NaN立刻減小學(xué)習(xí)率。更系統(tǒng)的做法是用學(xué)習(xí)率預(yù)熱和掃描比如在訓(xùn)練最初幾十個 step 里讓學(xué)習(xí)率從 0 線性增長到預(yù)設(shè)值觀察 loss 的下降曲線找到能穩(wěn)定訓(xùn)練的最大學(xué)習(xí)率。PyTorch 里設(shè)置學(xué)習(xí)率最常見的方式就是優(yōu)化器初始化時傳入optimizer torch.optim.SGD(model.parameters(), lr0.1)后續(xù)想動態(tài)調(diào)整學(xué)習(xí)率可以直接修改optimizer.param_groups[0][lr]不過更標準的方式是掛一個學(xué)習(xí)率調(diào)度器scheduler后面會詳細講。3.2 betasAdam 家族的“記憶長度”betas是 Adam 和 AdamW 里特有的參數(shù)是一個二元組。第一個值控制一階矩即動量的滑動平均衰減速度默認 0.9意思是每步更新時歷史動量的占比是 90%。第二個值控制二階矩的滑動平均衰減速度默認 0.999。你可以把 betas 理解為“記憶的窗口長度”。一階矩的窗口大約在1/(1-beta1)步即1/(1-0.9)10步二階矩的窗口大約是1/(1-0.999)1000步。這意味著 Adam 在估計梯度振幅時看的是過去大約 1000 步的梯度平方情況。什么情況下需要調(diào) betas常見的一種是訓(xùn)練 GAN 時有些人會把beta1調(diào)低到 0.5因為 GAN 的對抗訓(xùn)練需要更短的歷史記憶讓更新更激進一些。而beta2很少需要動除非你發(fā)現(xiàn)訓(xùn)練后期數(shù)值不穩(wěn)定可以嘗試把beta2調(diào)低到 0.99 或 0.995讓二階矩估計對近期梯度變化更敏感。optimizer torch.optim.Adam(model.parameters(), lr1e-3, betas(0.5, 0.999))3.3 weight_decay正則化與過擬合的博弈weight_decay在 PyTorch 里幾乎是所有優(yōu)化器都有的參數(shù)。它的作用是在每次參數(shù)更新后對所有權(quán)重乘一個略微小于 1 的系數(shù)從而抑制權(quán)重過大。這個操作等價于在損失函數(shù)中加入 L2 正則項但實現(xiàn)方式更直接。在 SGD 中weight_decay 與 L2 正則化是等價的。但在 Adam 里正如前面講的直接把 weight_decay 加進梯度會有問題所以 AdamW 才把它解耦出來。這也是我推薦大家用 AdamW 的原因之一。具體的 weight_decay 數(shù)值要看任務(wù)。常見范圍是0.0001到0.1之間。做圖像分類時1e-4是常見的起點做預(yù)訓(xùn)練語言模型微調(diào)時0.01甚至0.1都有使用。我在做 Transformer 模型時通常會從0.01起步如果驗證集過擬合明顯就往上調(diào)。3.4 eps一個容易被忽略卻保命的參數(shù)eps的作用是在計算分母時加一個極小值避免除零或極端數(shù)值導(dǎo)致的 NaN。它看似無關(guān)緊要但我確實碰到過因為eps問題導(dǎo)致訓(xùn)練崩潰的情況。默認eps1e-8對大部分模型都沒問題但在使用混合精度訓(xùn)練AMP時梯度幅值可能會被壓縮到很小1e-8這個量級的分母相對過高會影響自適應(yīng)學(xué)習(xí)率的精度。一些混合精度訓(xùn)練框架會把eps調(diào)整到1e-6或1e-7。如果你在 AMP 訓(xùn)練中碰到 loss 突然變成 NaN又排除了學(xué)習(xí)率太大的原因可以檢查一下eps是否合適。3.5 其他值得了解的參數(shù)amsgrad是 Adam 的一個小變體它保存了歷史上最大的二階矩估計從而避免學(xué)習(xí)率在某些情況下被人為增大。PyTorch 里通過amsgradTrue開啟。我在訓(xùn)練不穩(wěn)定的小批量任務(wù)時用過它有時能帶來穩(wěn)定性的提升但大多數(shù)情況下默認關(guān)閉即可。foreach是 PyTorch 2.x 新引入的性能優(yōu)化參數(shù)。設(shè)置為True時優(yōu)化器會使用更高效的遍歷方式更新參數(shù)在參數(shù)數(shù)量多、GPU 支持的情況下可以顯著提升速度。PyTorch 新版本里往往已經(jīng)默認開啟或自動選擇不需要特意管它。maximize參數(shù)用于做梯度上升而不是下降。某些對抗攻擊場景或者強化學(xué)習(xí)里需要最大化某個獎勵時可以把maximizeTrue這樣優(yōu)化器自動執(zhí)行梯度上升而不是下降。這個參數(shù)很多人不知道但用到的時候能省很多事。4. 優(yōu)化器論文里的干貨原理、缺陷與改進動機4.1 必讀的幾篇原文從 SGD 到 AdamW優(yōu)化器的論文大多是 2010 年代中期的經(jīng)典之作篇幅不長、公式清晰花一個下午讀原文比看十篇博客都有用。SGD 加動量的思想可以追溯到 Polyak 在 1964 年的工作但深度學(xué)習(xí)語境下更常引用的是 Sutskever 等人在 2013 年的論文《On the importance of initialization and momentum in deep learning》。這篇論文詳細分析了動量項如何加速收斂以及 Nesterov 動量的優(yōu)勢。Adam 的原文是 Kingma 和 Ba 的《Adam: A Method for Stochastic Optimization》2015 年發(fā)表。這篇論文里不僅有算法偽代碼還解答了一個關(guān)鍵問題為什么需要對一階矩和二階矩做偏差校正。因為指數(shù)滑動平均在初始時刻會偏向零偏差校正讓前幾步的估計更準確。PyTorch 的 Adam 實現(xiàn)里就包含了這個偏差校正邏輯。AdamW 對應(yīng)的論文是 Loshchilov 和 Hutter 的《Decoupled Weight Decay Regularization》。這篇論文理清了 L2 正則化與權(quán)重衰減在自適應(yīng)優(yōu)化器中的差別并通過大量實驗證明解耦權(quán)重衰減能取得更好的泛化性能。如果你想知道為什么現(xiàn)在大家都在用 AdamW這篇論文是最好的答案。還有一篇值得讀的是《On the Convergence of Adam and Beyond》它就是amsgrad的來源。這篇論文指出了 Adam 在某些情況下可能不收斂的問題并提出了一種修正策略。雖然 Amsgrad 在實際使用中提升有限但論文里對 Adam 收斂性的分析很能幫助你理解 Adam 的工作原理。4.2 論文里怎么報優(yōu)化器參數(shù)復(fù)現(xiàn)實驗必看我經(jīng)常需要復(fù)現(xiàn)別人論文里的實驗最頭疼的事情之一就是論文里優(yōu)化器配置交代得不清不楚。有些論文只寫了“我們使用 Adam學(xué)習(xí)率 1e-3”完全沒有提及 batch size、weight decay、學(xué)習(xí)率調(diào)度策略。而實際上這些細節(jié)往往嚴重影響結(jié)果。那些做得好的論文一般會在實驗設(shè)置部分明確寫出優(yōu)化器名稱、初始學(xué)習(xí)率、weight decay、總訓(xùn)練輪數(shù)和學(xué)習(xí)率衰減策略。比如經(jīng)典的 SwAV 自監(jiān)督學(xué)習(xí)論文里就會寫清楚使用了 LARS 優(yōu)化器、學(xué)習(xí)率怎么隨 batch size 線性縮放、cosine schedule 的細節(jié)。復(fù)現(xiàn)時要特別注意一個坑不同代碼庫對weight_decay的默認值不同。PyTorch 里SGD默認weight_decay0但有些論文的代碼實現(xiàn)里默認weight_decay5e-4。如果你照著論文里寫的超參數(shù)復(fù)現(xiàn)卻沒有注意到代碼庫的默認值差異訓(xùn)練結(jié)果可能天差地別。我一般復(fù)現(xiàn)時都會先打印一遍優(yōu)化器的param_groups把所有參數(shù)顯式寫出來避免默認值陷阱。4.3 從論文年代看優(yōu)化器的發(fā)展脈絡(luò)把優(yōu)化器論文按時間線排一下能看出深度學(xué)習(xí)訓(xùn)練方法論的演進。早期2012-2015是 SGD 和動量的天下。AlexNet、VGG、ResNet 這些經(jīng)典 CNN 模型的訓(xùn)練都是靠 SGD 加動量配合手動設(shè)計的學(xué)習(xí)率衰減策略。那個時代訓(xùn)練一個模型需要非常細致的調(diào)參經(jīng)驗學(xué)習(xí)率什么時候衰減、衰減多少倍都靠人對 loss 曲線的觀察判斷。中期2015-2018是自適應(yīng)方法爆發(fā)。Adagrad、RMSprop、Adam 相繼出現(xiàn)尤其是 Adam 以“開箱即用”的特性迅速成為主流。但很快研究者發(fā)現(xiàn) Adam 在泛化性能上不如 SGD于是開始深入分析原因推動了 AdamW 的誕生。近五年2019 至今的趨勢是“大模型 AdamW 固定學(xué)習(xí)率調(diào)度”。BERT、GPT、ViT 這些大規(guī)模 Transformer 模型的訓(xùn)練幾乎都是由 AdamW 主導(dǎo)的。究其原因Transformer 結(jié)構(gòu)對學(xué)習(xí)率極其敏感自適應(yīng)優(yōu)化器能夠更好地處理不同層之間的梯度尺度差異而解耦權(quán)重衰減又提供了更好的正則化效果。加上 warmup cosine decay 的調(diào)度策略形成了一整套相對成熟的訓(xùn)練范式。5. 實操PyTorch 優(yōu)化器的完整使用姿勢5.1 基礎(chǔ)訓(xùn)練循環(huán)從定義到 step這里給一個標準的 PyTorch 訓(xùn)練循環(huán)模板涵蓋優(yōu)化器、損失函數(shù)、反向傳播和參數(shù)更新。import torch import torch.nn as nn model nn.Linear(10, 2) optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay0.01) criterion nn.CrossEntropyLoss() for epoch in range(num_epochs): for batch_x, batch_y in dataloader: optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step()有幾個細節(jié)值得強調(diào)。optimizer.zero_grad()必須放在每次backward()之前因為 PyTorch 默認是累積梯度的。如果忘記清零上一個 batch 的梯度會累加到當(dāng)前 batch導(dǎo)致參數(shù)更新方向完全錯誤。很多人第一次寫訓(xùn)練循環(huán)就栽在這里。loss.backward()計算梯度optimizer.step()用梯度更新參數(shù)。這個順序千萬別反。有些新手會把optimizer.step()放在loss.backward()前梯度都還沒算出來就更新了參數(shù)自然一動不動。還有個習(xí)慣我要特別推薦在訓(xùn)練早期打印一次optimizer.state_dict()。這一下就能看到優(yōu)化器到底接收了哪些參數(shù)、每個參數(shù)組的學(xué)習(xí)率是多少、狀態(tài)是否為空。排查配置問題的時候非常高效。5.2 參數(shù)分組凍結(jié)與差異化學(xué)習(xí)率實際項目中很少會無腦優(yōu)化所有參數(shù)。最常見的情況是遷移學(xué)習(xí)用預(yù)訓(xùn)練模型做 backbone微調(diào)后面的分類頭。此時我們希望 backbone 參數(shù)變化幅度小分類頭參數(shù)從頭訓(xùn)起。PyTorch 的參數(shù)分組就是用param_groups實現(xiàn)差異化配置optimizer torch.optim.SGD([ {params: backbone.parameters(), lr: 1e-4}, {params: head.parameters(), lr: 1e-3}, ], lr1e-3, momentum0.9)注意這里的用法傳入的是一個字典列表每個字典表示一個參數(shù)組。默認lr寫在第二個參數(shù)位置被第一個參數(shù)組里的lr: 1e-4覆蓋。也就是說backbone 參數(shù)實際學(xué)習(xí)率是 1e-4head 參數(shù)使用默認的 1e-3其他未在字典里指定的參數(shù)組使用全局默認值。另一個常見需求是凍結(jié)某些層的參數(shù)。最簡單的辦法是設(shè)置param.requires_grad False但注意如果這些參數(shù)已經(jīng)被傳入優(yōu)化器優(yōu)化器里仍然會維護它們只是梯度為零、不更新。更高效的做法是把凍結(jié)參數(shù)從優(yōu)化器中剔除只傳入需要訓(xùn)練的參數(shù)。trainable_params [p for p in model.parameters() if p.requires_grad] optimizer torch.optim.Adam(trainable_params, lr1e-3)這樣參數(shù)少了很多優(yōu)化器每次 step 的計算開銷也會降低。5.3 學(xué)習(xí)率調(diào)度別讓學(xué)習(xí)率一成不變固定學(xué)習(xí)率訓(xùn)練到底層模型通常不是最優(yōu)選擇。經(jīng)驗上訓(xùn)練初期學(xué)習(xí)率要大一些快速下降后期學(xué)習(xí)率要小一些精細微調(diào)。PyTorch 的torch.optim.lr_scheduler提供了多種調(diào)度策略。最常用的是StepLR和CosineAnnealingLR。StepLR是每隔固定輪數(shù)把學(xué)習(xí)率乘一個衰減系數(shù)CosineAnnealingLR則讓學(xué)習(xí)率按照余弦曲線從初始值降到最小值。我現(xiàn)在做視覺模型偏愛 cosine因為它的下降曲線更平滑不需要手動設(shè)定衰減節(jié)點。scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxnum_epochs, eta_min1e-6 ) for epoch in range(num_epochs): train_one_epoch() scheduler.step()還有一個極其重要的策略是 warmup尤其在訓(xùn)練 Transformer 類模型時幾乎是必需項。warmup 的思想是訓(xùn)練最開始讓學(xué)習(xí)率從極小值線性增長到預(yù)設(shè)值避免模型初始化不穩(wěn)定時直接用大學(xué)習(xí)率導(dǎo)致震蕩。PyTorch 里可以用LambdaLR手寫def warmup_lambda(step): if step warmup_steps: return step / warmup_steps return 1.0 scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambdawarmup_lambda)5.4 梯度裁剪、AMP 與優(yōu)化器的配合梯度裁剪是防止梯度爆炸的常用手段在 RNN、Transformer、GAN 訓(xùn)練里尤其重要。PyTorch 提供了torch.nn.utils.clip_grad_norm_和clip_grad_value_兩個函數(shù)。區(qū)別是前者把梯度的范數(shù)限制在給定閾值內(nèi)后者把每個梯度值裁剪到給定范圍。loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step()順序有講究裁剪必須在backward()之后、optimizer.step()之前。因為這個操作修改的是參數(shù)梯度而 step 恰好要用到它。混合精度訓(xùn)練AMP如今在大模型訓(xùn)練中幾乎是標配。PyTorch 1.6 之后提供torch.cuda.amp接口。使用 AMP 時優(yōu)化器本身不需要大改但需要特別注意如果某些參數(shù)是float64或整數(shù)類型AMP 可能會報類型不匹配的錯誤。通常浮點參數(shù)用float32就沒問題。scaler torch.cuda.amp.GradScaler() for batch_x, batch_y in dataloader: optimizer.zero_grad() with torch.cuda.amp.autocast(): output model(batch_x) loss criterion(output, batch_y) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update()注意 AMP 訓(xùn)練時 grad clip 存在一個順序問題scaler.unscale_(optimizer)之前梯度是被放大過的直接用 clip 會導(dǎo)致閾值失真。上面的寫法先 unscale 再 clip 才是正確的。5.5 狀態(tài)保存與恢復(fù)斷點續(xù)訓(xùn)的完整姿勢訓(xùn)練大模型耗時動輒幾十個小時一旦中途中斷就得從頭開始這是不能接受的。好在 PyTorch 提供了優(yōu)化器狀態(tài)保存機制。torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, best_acc: best_acc, }, checkpoint.pt) # 恢復(fù) checkpoint torch.load(checkpoint.pt) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) epoch checkpoint[epoch]恢復(fù)優(yōu)化器狀態(tài)時有一點要特別注意必須在加載模型權(quán)重、并且在優(yōu)化器初始化之后才能調(diào)load_state_dict。而且如果模型結(jié)構(gòu)變了比如改了分類頭的維度優(yōu)化器里對應(yīng)的參數(shù)狀態(tài)會不匹配此時不能直接恢復(fù)。我一般會在結(jié)構(gòu)改動后重建優(yōu)化器不恢復(fù)優(yōu)化器狀態(tài)只是從斷點輪次重訓(xùn)。6. 常見問題與排查技巧實錄6.1 loss 不降反升到底是誰的鍋訓(xùn)練時最讓人崩潰的現(xiàn)象就是 loss 不但不降反而一路上漲。遇到這種情況我有一套固定的排查順序。先檢查學(xué)習(xí)率是不是太大。最簡單的驗證方法把學(xué)習(xí)率直接減小 10 倍如果 loss 曲線變得穩(wěn)定下降那問題基本就是學(xué)習(xí)率過高。再檢查輸入數(shù)據(jù)有沒有歸一化。不少新手直接拿原始像素 0-255 的圖片喂給模型梯度幅值巨大優(yōu)化器再強也很難穩(wěn)住。這時候損失函數(shù)干瞪眼不一定是優(yōu)化器的錯。最后檢查梯度是否出現(xiàn) NaN?？梢栽趏ptimizer.step()前打印torch.isnan(model.weight.grad).any()。如果有 NaN說明反向傳播里出了問題跟優(yōu)化器關(guān)系不大了。6.2 模型不收斂試試這些調(diào)參順序我調(diào)優(yōu)化器參數(shù)有個固定的優(yōu)先級順序。第一優(yōu)先看學(xué)習(xí)率因為它的影響最大。從論文或類似任務(wù)的默認值出發(fā)如果 loss 不降就降低學(xué)習(xí)率如果下降過慢就試著調(diào)高。第二是檢查 batch size 與學(xué)習(xí)率的匹配關(guān)系。batch size 翻倍學(xué)習(xí)率通常也該翻倍線性縮放原則。第三才是調(diào) weight_decay、momentum 這些細節(jié)。還有一個我特別推薦的做法先跑幾十個 step畫 loss 曲線觀察它是下降、震蕩還是發(fā)散。很多問題在早期就暴露了。沒必要一開始就跑完整訓(xùn)練流程。調(diào)參有一個反向直覺如果模型一直不收斂有時候不是優(yōu)化器調(diào)得不夠好而是模型結(jié)構(gòu)或數(shù)據(jù)預(yù)處理有問題。優(yōu)化器只能讓模型在損失面上運動如果損失面本身設(shè)計得不好再好的優(yōu)化器也白搭。6.3 優(yōu)化器相關(guān)報錯速查表報錯信息或現(xiàn)象常見原因解決方法ValueError: could not broadcast input array優(yōu)化器參數(shù)與模型參數(shù)維度不匹配檢查模型是否有參數(shù)被修改或新增重建優(yōu)化器RuntimeError: element 0 of tensors does not require grad傳入優(yōu)化器的參數(shù)沒有梯度設(shè)置requires_gradTrue或檢查模型是否在torch.no_grad內(nèi)Optimizer got an empty parameter list參數(shù)序列里沒有任何requires_gradTrue的參數(shù)檢查requires_grad設(shè)置或trainable_params篩選邏輯AMP 訓(xùn)練時GradScaler報錯優(yōu)化器與梯度縮放器配合次序錯誤按官方文檔順序scale(loss) - backward - unscale - stepload_state_dict時 size mismatch模型結(jié)構(gòu)改動導(dǎo)致優(yōu)化器狀態(tài)無法對齊不恢復(fù)優(yōu)化器狀態(tài)重新創(chuàng)建優(yōu)化器訓(xùn)練中途 loss 變成 NaN學(xué)習(xí)率過大、eps 過小、梯度爆炸先降學(xué)習(xí)率再檢查梯度范數(shù)必要時開 grad clip6.4 我的一個排查實例AdamW 微調(diào)模型驗證集掉點最后分享一個真實案例。我前段時間在拿一個預(yù)訓(xùn)練模型做下游任務(wù)微調(diào)一開始用 Adam、學(xué)習(xí)率 1e-4、無 weight decay驗證集分數(shù)一直上不去。后來我把優(yōu)化器換成 AdamW、weight_decay 設(shè)為 0.01、學(xué)習(xí)率調(diào)到 5e-5再配合 warmup驗證集分數(shù)明顯上升。整個過程優(yōu)化器參數(shù)只改了幾個數(shù)模型結(jié)構(gòu)完全沒動。這個例子很有代表性。它說明一個常被忽略的道理當(dāng)你覺得模型效果不好時不要只盯著網(wǎng)絡(luò)結(jié)構(gòu)改先回過頭檢查優(yōu)化器和學(xué)習(xí)率調(diào)度是否合理。很多時候優(yōu)化器配置的改變帶來的提升比換一個更大模型還要明顯。尤其是微調(diào)預(yù)訓(xùn)練模型學(xué)習(xí)率太大容易災(zāi)難性遺忘太小又學(xué)不到任務(wù)特有信息這個平衡點全靠優(yōu)化器配合調(diào)度策略來把握。在 PyTorch 里優(yōu)化器不是一個需要花很多時間去“學(xué)習(xí)”的模塊但絕對是一個值得花時間去“理解”的模塊。把torch.optim里每個常用優(yōu)化器的原理、參數(shù)、適用場景都摸清你在訓(xùn)練任何模型時都會更有底氣。剩下的就是不斷試錯、積累自己的調(diào)參感覺這本身就是深度學(xué)習(xí)實踐里最有趣的修煉過程。