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