1. 為什么“PyTorch”這三個(gè)字值得單獨(dú)拎出來(lái)聊如果你在深度學(xué)習(xí)圈子里待過(guò)一陣子大概率會(huì)發(fā)現(xiàn)一個(gè)現(xiàn)象大家聊模型結(jié)構(gòu)、聊訓(xùn)練技巧、聊部署方案的時(shí)候繞來(lái)繞去最后總會(huì)落回到一個(gè)詞上——PyTorch。它不是一個(gè)具體的項(xiàng)目名也不是某個(gè)新出的框架版本號(hào)而是一整套生態(tài)、一種寫(xiě)代碼的習(xí)慣、甚至可以說(shuō)是一種“思維方式”。我見(jiàn)過(guò)不少剛?cè)胄械呐笥焉蟻?lái)就問(wèn)“我該學(xué)TensorFlow還是PyTorch”也見(jiàn)過(guò)工作三五年的老手在重構(gòu)訓(xùn)練管線時(shí)毫不猶豫地把整個(gè)代碼庫(kù)遷到PyTorch上。這個(gè)標(biāo)題看起來(lái)簡(jiǎn)單到只有六個(gè)字母但它背后牽扯的東西非常多動(dòng)態(tài)圖與靜態(tài)圖的取舍、科研與生產(chǎn)的邊界、生態(tài)工具鏈的完整度、以及一個(gè)開(kāi)源項(xiàng)目如何從學(xué)術(shù)界一路殺到工業(yè)界主流。這篇文章我想做的事情很明確把PyTorch這個(gè)主題拆開(kāi)揉碎從設(shè)計(jì)哲學(xué)講到實(shí)操細(xì)節(jié)從環(huán)境配置講到訓(xùn)練調(diào)優(yōu)從常見(jiàn)坑點(diǎn)講到排查思路。不管你是剛裝完CUDA準(zhǔn)備跑第一個(gè)MNIST的新手還是已經(jīng)寫(xiě)過(guò)幾十個(gè)自定義算子、正在糾結(jié)分布式策略的老手我都盡量讓每一段都有能直接拿走用的東西。全文會(huì)圍繞PyTorch的核心機(jī)制、實(shí)操流程、工具選型、問(wèn)題排查這幾個(gè)維度展開(kāi)中間會(huì)穿插大量我在實(shí)際項(xiàng)目里踩過(guò)的坑和總結(jié)出來(lái)的經(jīng)驗(yàn)。文章不會(huì)涉及任何敏感話題所有案例均使用虛構(gòu)代稱(chēng)所有代碼和參數(shù)都基于公開(kāi)的通用實(shí)踐。先說(shuō)結(jié)論P(yáng)yTorch之所以能成為主流核心在于它把“定義即運(yùn)行”這件事做到了極致讓寫(xiě)模型像寫(xiě)普通Python腳本一樣自然。但自然不等于簡(jiǎn)單真正要把PyTorch用好你需要理解它的計(jì)算圖機(jī)制、內(nèi)存管理策略、數(shù)據(jù)加載管線、分布式通信原語(yǔ)以及一整套圍繞訓(xùn)練和部署的周邊工具。下面我會(huì)按照從整體設(shè)計(jì)到具體實(shí)現(xiàn)的順序一層一層往下拆。2. PyTorch整體設(shè)計(jì)思路與核心機(jī)制拆解2.1 動(dòng)態(tài)計(jì)算圖到底解決了什么問(wèn)題PyTorch最核心的設(shè)計(jì)選擇就是動(dòng)態(tài)計(jì)算圖也叫“定義即運(yùn)行”。什么意思呢你可以把它想象成做菜。靜態(tài)圖框架像是先把整份菜譜從頭到尾寫(xiě)死包括每一步用什么鍋、火候多大、什么時(shí)候放鹽全部確定之后才開(kāi)火。而動(dòng)態(tài)圖是你一邊切菜一邊嘗味道覺(jué)得淡了再加點(diǎn)鹽覺(jué)得火大了就調(diào)小一點(diǎn)。反映到代碼上靜態(tài)圖需要你先定義占位符、構(gòu)建計(jì)算圖、編譯、再喂數(shù)據(jù)動(dòng)態(tài)圖則是你寫(xiě)一行Python它就執(zhí)行一行計(jì)算圖是在運(yùn)行過(guò)程中動(dòng)態(tài)生成的。這個(gè)差異帶來(lái)的影響非常直接。第一調(diào)試變得極其自然。你可以在模型的前向傳播里隨便加print、加pdb.set_trace()、加條件判斷甚至根據(jù)中間結(jié)果動(dòng)態(tài)改變網(wǎng)絡(luò)結(jié)構(gòu)。第二控制流不再需要特殊語(yǔ)法。比如你要實(shí)現(xiàn)一個(gè)遞歸網(wǎng)絡(luò)或者一個(gè)根據(jù)輸入長(zhǎng)度決定循環(huán)次數(shù)的模塊在動(dòng)態(tài)圖里就是普通的for循環(huán)和if語(yǔ)句不需要引入tf.while_loop或者tf.cond這類(lèi)專(zhuān)用算子。第三科研迭代速度大幅提升。你改一個(gè)想法保存運(yùn)行立刻看到結(jié)果不需要等待漫長(zhǎng)的圖編譯過(guò)程。但動(dòng)態(tài)圖也不是沒(méi)有代價(jià)。早期的動(dòng)態(tài)圖框架在部署時(shí)面臨挑戰(zhàn)因?yàn)橛?jì)算圖是運(yùn)行時(shí)才確定的導(dǎo)出和優(yōu)化不如靜態(tài)圖方便。PyTorch的應(yīng)對(duì)策略是TorchScript和后來(lái)的torch.compile前者通過(guò)追蹤或腳本化把動(dòng)態(tài)代碼轉(zhuǎn)成靜態(tài)圖表示后者利用編譯技術(shù)對(duì)動(dòng)態(tài)圖進(jìn)行即時(shí)優(yōu)化。這個(gè)演進(jìn)過(guò)程本身就很值得琢磨它說(shuō)明PyTorch團(tuán)隊(duì)并沒(méi)有死守“純動(dòng)態(tài)”的教條而是在保留開(kāi)發(fā)體驗(yàn)的前提下逐步補(bǔ)齊生產(chǎn)部署的短板。2.2 張量抽象與自動(dòng)求導(dǎo)的配合邏輯PyTorch的另一塊基石是Tensor和autograd的緊密配合。Tensor本質(zhì)上是一個(gè)多維數(shù)組但它比NumPy數(shù)組多了兩樣?xùn)|西設(shè)備信息CPU還是GPU和梯度追蹤開(kāi)關(guān)requires_grad。當(dāng)你把一個(gè)張量的requires_grad設(shè)為T(mén)ruePyTorch就會(huì)開(kāi)始記錄所有對(duì)它進(jìn)行的操作構(gòu)建一張反向傳播用的計(jì)算圖。每個(gè)操作都會(huì)生成一個(gè)Function節(jié)點(diǎn)節(jié)點(diǎn)里保存了前向傳播的輸入和反向傳播需要的中間變量。這里有一個(gè)很容易被忽略的細(xì)節(jié)計(jì)算圖是在前向傳播過(guò)程中動(dòng)態(tài)構(gòu)建的每次迭代都會(huì)重新構(gòu)建。這意味著你不能像靜態(tài)圖那樣“編譯一次、運(yùn)行多次”但也意味著你可以根據(jù)每次迭代的不同情況靈活調(diào)整計(jì)算路徑。比如在強(qiáng)化學(xué)習(xí)里你可能會(huì)根據(jù)當(dāng)前策略決定是否執(zhí)行某個(gè)分支動(dòng)態(tài)圖天然支持這種模式。自動(dòng)求導(dǎo)的觸發(fā)時(shí)機(jī)也很關(guān)鍵。默認(rèn)情況下只有標(biāo)量才能直接調(diào)用.backward()因?yàn)镻yTorch需要從標(biāo)量開(kāi)始反向傳播。如果你有一個(gè)多維輸出需要先通過(guò)損失函數(shù)把它聚合成標(biāo)量或者傳入一個(gè)梯度參數(shù)。這個(gè)設(shè)計(jì)是為了避免歧義多維輸出的反向傳播需要指定每個(gè)輸出維度的梯度權(quán)重PyTorch選擇讓你顯式處理這件事。2.3 模塊化設(shè)計(jì)與生態(tài)分層PyTorch的代碼組織非常清晰大致可以分為幾層。最底層是torch核心庫(kù)提供張量運(yùn)算和自動(dòng)求導(dǎo)。往上是torch.nn提供神經(jīng)網(wǎng)絡(luò)層、損失函數(shù)、容器等抽象。再往上是torch.optim提供各種優(yōu)化器。然后是torch.utils.data負(fù)責(zé)數(shù)據(jù)加載和預(yù)處理。最外層是領(lǐng)域?qū)Ｓ脦?kù)比如torchvision處理圖像、torchaudio處理音頻、torchtext處理文本。這種分層設(shè)計(jì)的好處是每一層都可以獨(dú)立使用。你可以只用torch.Tensor做數(shù)值計(jì)算完全不用nn.Module也可以只用nn.Module組織模型自己寫(xiě)訓(xùn)練循環(huán)而不依賴optim。這種靈活性是PyTorch吸引科研人員的重要原因之一。但靈活性也意味著你需要自己決定用哪一層抽象對(duì)于新手來(lái)說(shuō)選擇太多反而容易迷失。我的建議是剛開(kāi)始的時(shí)候盡量用高層抽象比如nn.Sequential、DataLoader、optim.Adam先把流程跑通。等你對(duì)整體流程有了感覺(jué)再逐步往下深入理解每一層到底做了什么。不要一上來(lái)就手寫(xiě)自定義Function或者自己實(shí)現(xiàn)優(yōu)化器那樣很容易在細(xì)節(jié)里迷失方向。3. 環(huán)境配置與第一個(gè)訓(xùn)練管線的實(shí)操要點(diǎn)3.1 安裝策略與版本匹配的坑PyTorch的安裝看起來(lái)簡(jiǎn)單一行pip install torch就能搞定但實(shí)際項(xiàng)目里版本匹配是個(gè)大問(wèn)題。你需要同時(shí)考慮Python版本、CUDA版本、cuDNN版本、PyTorch版本以及各種周邊庫(kù)的兼容性。我見(jiàn)過(guò)太多人因?yàn)榘姹静黄ヅ鋵?dǎo)致ImportError或者運(yùn)行時(shí)崩潰浪費(fèi)半天時(shí)間。一個(gè)比較穩(wěn)妥的做法是先確定你的顯卡驅(qū)動(dòng)支持的CUDA版本然后去PyTorch官網(wǎng)的安裝頁(yè)面選擇對(duì)應(yīng)的命令。不要盲目用pip install torch因?yàn)槟J(rèn)安裝的可能是CPU版本或者CUDA版本和你的驅(qū)動(dòng)不匹配。如果你用conda可以用conda install pytorch torchvision torchaudio cudatoolkit11.8 -c pytorch這樣的命令conda會(huì)自動(dòng)處理依賴關(guān)系。另一個(gè)坑是虛擬環(huán)境。強(qiáng)烈建議每個(gè)項(xiàng)目單獨(dú)建一個(gè)虛擬環(huán)境用venv或者conda都行。我自己的習(xí)慣是用conda因?yàn)樗谔幚鞢UDA相關(guān)依賴時(shí)更省心。建好環(huán)境后先裝PyTorch再裝其他庫(kù)避免其他庫(kù)把PyTorch的依賴降級(jí)或升級(jí)。注意如果你在Windows上開(kāi)發(fā)CUDA版本和PyTorch版本的對(duì)應(yīng)關(guān)系比Linux更嚴(yán)格建議直接去官網(wǎng)查兼容性表格不要憑經(jīng)驗(yàn)猜。3.2 數(shù)據(jù)加載管線的設(shè)計(jì)細(xì)節(jié)數(shù)據(jù)加載是訓(xùn)練管線里最容易被低估的部分。很多人把注意力全放在模型結(jié)構(gòu)上結(jié)果訓(xùn)練速度慢得離譜最后發(fā)現(xiàn)瓶頸在數(shù)據(jù)加載。PyTorch的DataLoader提供了多進(jìn)程加載、批處理、打亂、預(yù)取等功能但默認(rèn)參數(shù)不一定適合你的場(chǎng)景。num_workers是最關(guān)鍵的參數(shù)之一。它決定了用多少個(gè)子進(jìn)程來(lái)加載數(shù)據(jù)。設(shè)得太小GPU等數(shù)據(jù)設(shè)得太大CPU和內(nèi)存壓力大甚至可能因?yàn)檫M(jìn)程間通信開(kāi)銷(xiāo)反而變慢。我的經(jīng)驗(yàn)是從num_workers4開(kāi)始試觀察GPU利用率。如果GPU利用率經(jīng)常掉到50%以下說(shuō)明數(shù)據(jù)加載是瓶頸可以適當(dāng)增加如果CPU占用已經(jīng)很高再增加也沒(méi)用。pin_memory是另一個(gè)重要參數(shù)。當(dāng)你在GPU上訓(xùn)練時(shí)把pin_memoryTrue可以讓數(shù)據(jù)加載器把張量放在鎖頁(yè)內(nèi)存里這樣從CPU傳到GPU的速度會(huì)更快。這個(gè)參數(shù)對(duì)訓(xùn)練速度的影響有時(shí)候能達(dá)到10%以上但很多人不知道。batch_size的選擇需要綜合考慮顯存大小、模型復(fù)雜度和收斂速度。一般來(lái)說(shuō)較大的batch size能提高GPU利用率但可能影響收斂質(zhì)量。一個(gè)常用的策略是先用較大的batch size跑通流程確認(rèn)沒(méi)有顯存問(wèn)題后再根據(jù)收斂情況調(diào)整。如果顯存不夠可以用梯度累積來(lái)模擬大batch的效果。3.3 訓(xùn)練循環(huán)的骨架與關(guān)鍵控制點(diǎn)一個(gè)標(biāo)準(zhǔn)的PyTorch訓(xùn)練循環(huán)包含以下幾個(gè)步驟前向傳播、計(jì)算損失、清空梯度、反向傳播、更新參數(shù)。看起來(lái)簡(jiǎn)單但每個(gè)步驟都有細(xì)節(jié)。清空梯度這一步經(jīng)常被忘記。PyTorch默認(rèn)會(huì)累積梯度如果你不在每次迭代開(kāi)始時(shí)調(diào)用optimizer.zero_grad()梯度會(huì)一直累加導(dǎo)致更新方向錯(cuò)誤。我見(jiàn)過(guò)不少新手因?yàn)檫@個(gè)原因訓(xùn)練不收斂排查半天才發(fā)現(xiàn)是漏了這一行。反向傳播之后、更新參數(shù)之前有時(shí)候需要做梯度裁剪。特別是在訓(xùn)練RNN或者Transformer時(shí)梯度爆炸是常見(jiàn)問(wèn)題。torch.nn.utils.clip_grad_norm_可以按范數(shù)裁剪梯度clip_grad_value_可以按值裁剪。裁剪閾值需要根據(jù)具體任務(wù)調(diào)整一般從1.0開(kāi)始試。學(xué)習(xí)率調(diào)度也是訓(xùn)練循環(huán)的一部分。PyTorch提供了lr_scheduler模塊支持StepLR、CosineAnnealingLR、OneCycleLR等策略。我的習(xí)慣是在訓(xùn)練初期用warmup讓學(xué)習(xí)率從一個(gè)小值逐漸升到目標(biāo)值避免一開(kāi)始就大步長(zhǎng)導(dǎo)致不穩(wěn)定。warmup的步數(shù)一般是總步數(shù)的5%到10%。提示如果你發(fā)現(xiàn)loss突然變成NaN優(yōu)先檢查學(xué)習(xí)率是不是太大其次檢查數(shù)據(jù)里有沒(méi)有異常值最后檢查模型里有沒(méi)有數(shù)值不穩(wěn)定的操作比如log(0)或者除以零。4. 模型定義、調(diào)試與性能優(yōu)化的實(shí)戰(zhàn)經(jīng)驗(yàn)4.1 nn.Module的組織方式與常見(jiàn)誤區(qū)nn.Module是PyTorch組織模型的標(biāo)準(zhǔn)方式。一個(gè)典型的模塊包含__init__和forward兩個(gè)方法。__init__里定義子模塊和參數(shù)forward里定義計(jì)算邏輯?？雌饋?lái)簡(jiǎn)單但有幾個(gè)容易踩的坑。第一個(gè)坑是忘記調(diào)用super().__init__()。如果你在自定義模塊里不調(diào)用父類(lèi)的初始化方法PyTorch就無(wú)法正確注冊(cè)子模塊和參數(shù)后續(xù)的parameters()、to(device)、state_dict()都會(huì)出問(wèn)題。這個(gè)錯(cuò)誤在IDE里不一定報(bào)錯(cuò)但運(yùn)行時(shí)會(huì)出各種奇怪的問(wèn)題。第二個(gè)坑是把張量直接賦值為模塊屬性。比如self.some_tensor torch.randn(10)這樣賦值的張量不會(huì)被注冊(cè)為參數(shù)或緩沖區(qū)不會(huì)隨模型一起移動(dòng)到GPU也不會(huì)出現(xiàn)在state_dict里。如果你希望它隨模型移動(dòng)應(yīng)該用register_buffer如果希望它被訓(xùn)練應(yīng)該用nn.Parameter。第三個(gè)坑是在forward里創(chuàng)建新的模塊或參數(shù)。每次調(diào)用forward都創(chuàng)建新模塊會(huì)導(dǎo)致內(nèi)存泄漏和性能下降。正確的做法是在__init__里創(chuàng)建好在forward里復(fù)用。4.2 調(diào)試工具與技巧實(shí)錄PyTorch的調(diào)試體驗(yàn)是它的一大優(yōu)勢(shì)。你可以像調(diào)試普通Python代碼一樣調(diào)試模型。我常用的幾個(gè)技巧包括在forward里插入print語(yǔ)句查看張量形狀和數(shù)值范圍用torch.autograd.set_detect_anomaly(True)檢測(cè)反向傳播中的異常用torch.utils.bottleneck分析性能瓶頸用torch.profiler做更細(xì)粒度的性能分析。torch.profiler是我強(qiáng)烈推薦的工具。它可以記錄每個(gè)算子的執(zhí)行時(shí)間、內(nèi)存占用、調(diào)用棧等信息幫你快速定位性能瓶頸。用法也很簡(jiǎn)單用with torch.profiler.profile(...) as prof:包裹訓(xùn)練循環(huán)然后print(prof.key_averages().table(...))輸出結(jié)果。我一般會(huì)關(guān)注cuda_time_total和cpu_time_total兩列看看時(shí)間主要花在哪里。另一個(gè)實(shí)用技巧是用hooks查看中間層的輸出。register_forward_hook可以在前向傳播時(shí)捕獲某一層的輸出register_backward_hook可以捕獲反向傳播的梯度。這對(duì)于分析梯度消失或爆炸非常有用。4.3 混合精度訓(xùn)練與顯存優(yōu)化混合精度訓(xùn)練是提升訓(xùn)練速度和減少顯存占用的有效手段。PyTorch從1.6版本開(kāi)始內(nèi)置了torch.cuda.amp用起來(lái)很方便。核心思路是用autocast上下文管理器自動(dòng)把部分算子轉(zhuǎn)成半精度同時(shí)用GradScaler處理梯度縮放避免半精度下的下溢問(wèn)題。實(shí)際使用中有幾個(gè)細(xì)節(jié)需要注意。第一不是所有算子都適合半精度autocast會(huì)自動(dòng)判斷但你也可以手動(dòng)指定某些區(qū)域用全精度。第二GradScaler需要在反向傳播前調(diào)用scale(loss)在更新參數(shù)前調(diào)用step(optimizer)和update()。第三混合精度訓(xùn)練可能會(huì)影響收斂如果發(fā)現(xiàn)loss波動(dòng)變大可以嘗試調(diào)整學(xué)習(xí)率或者關(guān)閉混合精度。顯存優(yōu)化的另一個(gè)手段是梯度檢查點(diǎn)。torch.utils.checkpoint可以讓你用計(jì)算換顯存在前向傳播時(shí)不保存中間激活值反向傳播時(shí)重新計(jì)算。這對(duì)于訓(xùn)練大模型非常有用代價(jià)是訓(xùn)練速度會(huì)慢一些。使用的時(shí)候要注意檢查點(diǎn)只適用于nn.Module的forward方法而且需要確保輸入是可求導(dǎo)的。5. 分布式訓(xùn)練與部署上線的關(guān)鍵決策5.1 數(shù)據(jù)并行與模型并行的選擇邏輯當(dāng)單卡顯存不夠或者訓(xùn)練速度太慢時(shí)就需要考慮分布式訓(xùn)練。PyTorch提供了多種并行策略最常用的是DataParallel和DistributedDataParallel。DataParallel用起來(lái)簡(jiǎn)單一行代碼就能把模型放到多卡上但它有個(gè)致命缺點(diǎn)主卡負(fù)載高通信效率低而且不支持多機(jī)。所以生產(chǎn)環(huán)境基本都用DistributedDataParallel。DistributedDataParallel的核心思想是每個(gè)進(jìn)程控制一張卡各自加載數(shù)據(jù)、計(jì)算梯度然后通過(guò)all_reduce操作同步梯度。它的通信效率比DataParallel高很多而且支持多機(jī)多卡。但配置起來(lái)也復(fù)雜一些需要設(shè)置init_process_group、DistributedSampler、local_rank等。模型并行適用于單層參數(shù)就超過(guò)單卡顯存的情況比如超大Transformer。PyTorch提供了torch.distributed.pipeline.sync.Pipe和torch.distributed.tensor.parallel等工具但使用門(mén)檻較高需要手動(dòng)切分模型和通信。我的建議是優(yōu)先考慮數(shù)據(jù)并行只有在數(shù)據(jù)并行解決不了問(wèn)題的時(shí)候才上模型并行。5.2 模型導(dǎo)出與推理優(yōu)化訓(xùn)練完之后模型需要導(dǎo)出用于推理。PyTorch提供了多種導(dǎo)出方式torch.save保存整個(gè)模型或state_dicttorch.jit.trace或torch.jit.script導(dǎo)出TorchScripttorch.onnx.export導(dǎo)出ONNX格式。每種方式都有適用場(chǎng)景。state_dict是最輕量的保存方式只保存參數(shù)不保存模型結(jié)構(gòu)。加載的時(shí)候需要先實(shí)例化模型再load_state_dict。這種方式的好處是靈活壞處是模型結(jié)構(gòu)變了之后可能加載失敗。TorchScript適合需要脫離Python環(huán)境部署的場(chǎng)景。trace方式適合沒(méi)有控制流的模型script方式適合有控制流的模型。導(dǎo)出之后可以用torch.jit.load加載在C環(huán)境里運(yùn)行。ONNX適合跨框架部署比如把模型部署到TensorRT或者OpenVINO上。導(dǎo)出的時(shí)候需要注意算子兼容性有些PyTorch算子ONNX不支持需要替換或自定義。推理優(yōu)化方面torch.compile是最近比較火的方向。它通過(guò)即時(shí)編譯把動(dòng)態(tài)圖轉(zhuǎn)成優(yōu)化的靜態(tài)圖在推理和訓(xùn)練上都有加速效果。不過(guò)torch.compile對(duì)模型結(jié)構(gòu)有一定要求不是所有模型都能順利編譯。我的經(jīng)驗(yàn)是先跑通再優(yōu)化不要一上來(lái)就追求極致性能。5.3 常見(jiàn)部署問(wèn)題與排查思路部署階段最常見(jiàn)的問(wèn)題是推理結(jié)果和訓(xùn)練結(jié)果不一致。原因可能有很多預(yù)處理不一致、后處理不一致、數(shù)值精度差異、算子實(shí)現(xiàn)差異等。排查的時(shí)候我一般會(huì)先用相同的輸入分別跑訓(xùn)練模式和推理模式對(duì)比中間層的輸出逐步縮小范圍。另一個(gè)常見(jiàn)問(wèn)題是性能不達(dá)標(biāo)。如果推理速度比預(yù)期慢很多優(yōu)先檢查是不是用了CPU推理而沒(méi)用GPU或者是不是沒(méi)有開(kāi)torch.no_grad()導(dǎo)致構(gòu)建了不必要的計(jì)算圖。torch.no_grad()在推理時(shí)非常重要它不僅能減少顯存占用還能提升速度。還有一個(gè)坑是模型加載失敗。如果load_state_dict報(bào)錯(cuò)說(shuō)key不匹配可能是模型結(jié)構(gòu)變了或者保存的時(shí)候用了DataParallel導(dǎo)致key多了module.前綴。解決辦法是用model.module.load_state_dict或者手動(dòng)去掉前綴。6. 常見(jiàn)問(wèn)題速查與避坑經(jīng)驗(yàn)匯總6.1 訓(xùn)練不收斂的排查清單訓(xùn)練不收斂是最高頻的問(wèn)題之一。我整理了一個(gè)排查順序從簡(jiǎn)單到復(fù)雜先檢查數(shù)據(jù)有沒(méi)有問(wèn)題比如標(biāo)簽是不是對(duì)的、輸入有沒(méi)有歸一化、有沒(méi)有異常值再檢查損失函數(shù)對(duì)不對(duì)比如分類(lèi)任務(wù)用交叉熵、回歸任務(wù)用MSE然后檢查學(xué)習(xí)率是不是太大或太小可以嘗試用學(xué)習(xí)率掃描接著檢查模型結(jié)構(gòu)有沒(méi)有問(wèn)題比如維度不匹配、激活函數(shù)選錯(cuò)最后檢查梯度有沒(méi)有異常用torch.autograd.set_detect_anomaly(True)或者手動(dòng)打印梯度范數(shù)?，F(xiàn)象可能原因排查方法loss不下降學(xué)習(xí)率太小嘗試增大10倍loss震蕩學(xué)習(xí)率太大嘗試減小10倍loss變NaN梯度爆炸加梯度裁剪訓(xùn)練loss降但驗(yàn)證loss升過(guò)擬合加正則化或早停準(zhǔn)確率不變標(biāo)簽或數(shù)據(jù)有問(wèn)題檢查數(shù)據(jù)加載邏輯6.2 顯存不足的應(yīng)對(duì)策略顯存不足是另一個(gè)高頻問(wèn)題。應(yīng)對(duì)策略按優(yōu)先級(jí)排序第一減小batch size第二用混合精度訓(xùn)練第三用梯度檢查點(diǎn)第四清理不必要的緩存比如torch.cuda.empty_cache()第五檢查有沒(méi)有在訓(xùn)練循環(huán)里累積了不必要的張量比如把loss存到列表里但忘了detach。注意torch.cuda.empty_cache()只能釋放未被使用的緩存不能釋放正在使用的顯存。如果你發(fā)現(xiàn)調(diào)用之后顯存沒(méi)變化說(shuō)明顯存確實(shí)被占用了需要從代碼層面找原因。6.3 多卡訓(xùn)練中的同步問(wèn)題多卡訓(xùn)練時(shí)BatchNorm的同步是個(gè)容易被忽略的問(wèn)題。默認(rèn)情況下每張卡上的BatchNorm獨(dú)立計(jì)算均值和方差這可能導(dǎo)致訓(xùn)練不穩(wěn)定。PyTorch提供了SyncBatchNorm可以在多卡之間同步統(tǒng)計(jì)量。用法是model torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)在DistributedDataParallel之前調(diào)用。另一個(gè)同步問(wèn)題是隨機(jī)種子。多卡訓(xùn)練時(shí)每張卡的數(shù)據(jù)加載和隨機(jī)增強(qiáng)需要有不同的種子否則可能加載到相同的數(shù)據(jù)。DistributedSampler會(huì)自動(dòng)處理這個(gè)問(wèn)題但如果你自己寫(xiě)數(shù)據(jù)加載邏輯需要手動(dòng)設(shè)置epoch和rank相關(guān)的種子。6.4 版本升級(jí)帶來(lái)的兼容性坑PyTorch版本迭代很快每次大版本升級(jí)都可能引入不兼容的改動(dòng)。比如從1.x到2.xtorch.compile成為默認(rèn)推薦但一些舊的寫(xiě)法可能不再被支持。升級(jí)之前建議先看官方release note里的breaking changes然后在測(cè)試環(huán)境里跑一遍完整的訓(xùn)練和推理流程。我自己的習(xí)慣是生產(chǎn)環(huán)境鎖定版本不輕易升級(jí)實(shí)驗(yàn)環(huán)境可以嘗鮮但要做好回滾準(zhǔn)備。如果必須升級(jí)先用一個(gè)小數(shù)據(jù)集跑通全流程確認(rèn)沒(méi)有報(bào)錯(cuò)和性能下降再全量遷移。7. 我個(gè)人在實(shí)際操作中的幾點(diǎn)體會(huì)PyTorch用久了會(huì)發(fā)現(xiàn)它最大的價(jià)值不在于某個(gè)具體功能而在于它把“可調(diào)試性”和“靈活性”放在了很高的優(yōu)先級(jí)。這意味著你可以用很自然的方式表達(dá)復(fù)雜的想法而不用跟框架的抽象做斗爭(zhēng)。但這也意味著你需要自己對(duì)很多細(xì)節(jié)負(fù)責(zé)比如內(nèi)存管理、設(shè)備同步、梯度控制。框架不會(huì)幫你做所有事情它給你工具怎么用取決于你。我踩過(guò)最深的坑是在分布式訓(xùn)練里。一開(kāi)始覺(jué)得DistributedDataParallel配置太麻煩就用DataParallel湊合結(jié)果訓(xùn)練速度比單卡還慢。后來(lái)老老實(shí)實(shí)配DistributedDataParallel才發(fā)現(xiàn)多卡訓(xùn)練的加速比可以接近線性。所以我的建議是如果決定上多卡就直接用DistributedDataParallel不要走彎路。另一個(gè)體會(huì)是不要過(guò)早優(yōu)化。先把模型跑通確認(rèn)收斂再考慮混合精度、梯度檢查點(diǎn)、分布式這些高級(jí)特性。很多性能問(wèn)題在流程跑通之后自然就暴露出來(lái)了這時(shí)候再針對(duì)性優(yōu)化效率更高。最后分享一個(gè)小技巧如果你不確定某個(gè)操作會(huì)不會(huì)影響梯度可以用torch.autograd.gradcheck做數(shù)值梯度檢查。雖然慢但能幫你確認(rèn)自定義算子的反向傳播實(shí)現(xiàn)是否正確。這個(gè)工具在寫(xiě)自定義Function的時(shí)候特別有用。