總有人問我學 PyTorch 第一關難在哪。我的答案不是張量、不是模型結(jié)構(gòu)也不是那堆層 API而是把環(huán)境裝好。真的我見過太多人在配置環(huán)境這一步折騰到懷疑人生最后連import torch都沒跑過就直接放棄了。這篇就聊聊我自己的實踐路徑不講大而全的框架介紹只講一件事怎么少踩坑地把 PyTorch 環(huán)境搞定以及跑通第一個訓練腳本時真正需要懂的那幾個核心概念。內(nèi)容面向剛?cè)腴T、想快速跑起來的人也適合那些已經(jīng)能跑通 Demo、但對一些底層機制還一知半解的讀者。1. 裝環(huán)境這件事九成新手都卡在第一步1.1 先確認硬件再談安裝方式拿到一臺機器我建議做的第一件事就是看顯卡型號而不是先上網(wǎng)搜教程。PyTorch 的安裝方式完全取決于你有沒有 NVIDIA 顯卡、驅(qū)動支持什么版本的 CUDA這步搞錯了后面全是白折騰。Windows 下打開設備管理器看“顯示適配器”Linux 下跑一句nvidia-smi如果這命令不存在說明驅(qū)動都沒裝好。這里有個常見的誤區(qū)nvidia-smi顯示的是驅(qū)動支持的CUDA 版本上限并不是你系統(tǒng)里已經(jīng)裝好的 CUDA 版本。很多教程會讓人先去裝 CUDA Toolkit但用 PyTorch 官方預編譯包的話這一步基本可以跳過。你只需要一個足夠新的驅(qū)動剩下的 CUDA 運行庫都跟著 PyTorch 的 wheel 包一起走。如果機器沒有獨立顯卡或者買的是某果芯片的機器那就只能走 CPU 版本。CPU 版對學習和調(diào)通代碼完全夠用但訓練大規(guī)模模型就別想了這點需要有預期。1.2 CUDA 版本到底怎么選很多人在這卡住是因為官方安裝命令里給了 cuda 12.4、cuda 12.1、cuda 11.8 好幾個選項不知道選哪個。我的建議很直接看nvidia-smi頂部顯示的 Driver 版本對應的 CUDA 支持范圍然后選一個不低于 PyTorch 要求的版本號即可。舉個例子某種驅(qū)動版本支持的 CUDA 是 12.x那你就直接選官方給的 cuda 12.4 或更高版本對應的安裝命令。你不需要擔心選錯了會怎樣只要驅(qū)動版本不落后太多PyTorch 的依賴庫都能正常運行。實際上CUDA 自帶向下兼容能力只要驅(qū)動支持的 CUDA 版本高于 PyTorch 所需的副本次數(shù)就沒事。1.3 用 conda 還是 pip我的建議很折騰說句實話兩個都能用但如果你是在自己機器上做實驗我更傾向用 conda 建干凈的環(huán)境。原因是 conda 不僅管 Python 包還能管一部分系統(tǒng)依賴庫遇到坑時回滾方便。使用 pip 的話依賴關系經(jīng)常把系統(tǒng)搞亂。conda create -n torch_env python3.10 conda activate torch_env pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124這里有個細節(jié)即便在 conda 環(huán)境里我也建議最后用 pip 裝 PyTorch。因為 PyTorch 官方對 conda 頻道和 pip 兩套渠道的同步速度不一樣pip 裝的版往往比 conda 頻道更新更快而且用 pip 裝預編譯包還省事。命令末尾那個--index-url參數(shù)就是指定從官方庫的 CUDA 版本渠道下載比先裝 CUDA 再裝 torch 的方式要簡單得多。裝完第一件事跑一段驗證代碼import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))這三行能正確輸出說明環(huán)境已經(jīng)通了。記住torch.cuda.is_available()返回True才是重點版本號亂不亂無所謂。提示不要強迫癥一樣非要裝最新版。PyTorch 生態(tài)里有些經(jīng)典代碼庫可能還沒跟上新版 API選擇一個穩(wěn)定版本比如 2.1 或 2.2往往更省心。2. Tensor 和 autograd真正需要搞懂的兩個地基2.1 Tensor 并不神秘它就是帶“加速器”的多維數(shù)組Tensor 這個詞聽起來高大上但你完全可以把它理解為帶各種附屬信息的多維數(shù)組。它和 NumPy 的 ndarray 很像區(qū)別在于 Tensor 能自動在 GPU 上做運算還能記錄計算過程。舉個我常給朋友講的例子把 Tensor 想象成一張快遞單。它本身裝著數(shù)據(jù)好比包裹還附著一些“物流信息”——在哪個設備上CPU 還是 GPU、需要不需要記錄梯度requires_grad、梯度的值是什么grad。普通數(shù)組就像你自己拎著包裹走Tensor 則相當于走了一套完整的物流系統(tǒng)每一步都有跡可循。在代碼層面最常用的創(chuàng)建方式有幾種import torch # 從列表創(chuàng)建 x torch.tensor([1.0, 2.0, 3.0]) # 全零 / 全一 / 隨機 y torch.zeros(2, 3) z torch.ones(2, 3) w torch.randn(2, 3) # 像 NumPy 一樣做運算 result x.sum()Tensor 默認在 CPU 上想挪到 GPU 就調(diào)用.to(cuda)。不過別急著把每個數(shù)據(jù)都手動搬來搬去后面用DataLoader和.to(device)的模式會更規(guī)范。2.2 autograd 是怎么做到“不用手算梯度”的這是初學 PyTorch 時最值得多花時間理解的一塊。傳統(tǒng)方法里你要自己推導出每一層參數(shù)的梯度公式不僅繁瑣還容易出錯。PyTorch 的自動求導機制干的事是把你的計算過程記錄成一張圖然后按鏈式法則反過來算梯度?？梢园堰@個過程想象成順著水管走一遍再逆著水管走一遍。前向傳播的時候水從進口流向出口反向傳播的時候水流的方向反過來每個閥門參數(shù)都能知道自己被推了多少力。這就是為什么只要設置requires_gradTrue然后在 loss 上調(diào)用.backward()所有參數(shù)的.grad就自動填上了值。有個點很多人會忽略默認情況下Tensor 的requires_grad是False。只有當你打算對這個張量求梯度才需要設成True。對模型參數(shù)來說nn.Parameter會默認開啟這個屬性所以平時你不太需要手動管。再看一段簡單的代碼x torch.tensor([2.0], requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad) # 結(jié)果應該是 7因為 2x3 在 x2 時是 7這代碼雖然簡單但背后的意義不簡單——對學習率、反向傳播流程的理解全靠這個機制打底。2.3 用 autograd 跑通第一個最小示例不急著上模型庫先用最原始的方式手寫一層邏輯這樣理解最深。假設我們要擬合一條直線y w * x bimport torch w torch.randn(1, requires_gradTrue) b torch.zeros(1, requires_gradTrue) x_data torch.tensor([1.0, 2.0, 3.0, 4.0]) y_data torch.tensor([2.0, 4.0, 6.0, 8.0]) learning_rate 0.01 for epoch in range(100): y_pred w * x_data b loss ((y_pred - y_data) ** 2).mean() loss.backward() with torch.no_grad(): w - learning_rate * w.grad b - learning_rate * b.grad w.grad.zero_() b.grad.zero_() if epoch % 20 0: print(fepoch {epoch}: loss{loss.item():.4f})這個代碼里有兩個值得留意的操作一個是torch.no_grad()它告訴 PyTorch 這段操作不需要記錄計算圖因為參數(shù)的更新本身不該被納入梯度計算另一個是每次更新后要調(diào)zero_()清空梯度不然下一次backward()會把新舊梯度累加起來結(jié)果就錯了。把這段跑通你會對訓練的本質(zhì)有體感計算預測值、算 loss、反向傳播、更新參數(shù)、清零梯度這個過程循環(huán)往復。之后再用nn.Linear這些高層 API也只是把這個流程封裝了而已。3. 從玩具代碼到能跑的腳本數(shù)據(jù)管道和訓練流程的細節(jié)3.1 DataLoader別小看數(shù)據(jù)入口的寫法很多新手一開始圖省事把整批數(shù)據(jù)一次性加載進內(nèi)存然后手動切 batch。數(shù)據(jù)量小的時候沒什么感覺數(shù)據(jù)一多就直接內(nèi)存爆掉。PyTorch 提供了一套標準的數(shù)據(jù)處理流程核心是Dataset和DataLoader兩個類。你只需要定義好一個繼承自Dataset的類實現(xiàn)__len__和__getitem__兩個方法剩下的 batch 分割、亂序、多進程讀取都由DataLoader處理。這樣做的好處很明顯數(shù)據(jù)量大時不用一次全讀入內(nèi)存按需讀取即可。import torch from torch.utils.data import Dataset, DataLoader import torchvision.transforms as T class MyDataset(Dataset): def __init__(self, data, targets): self.data data self.targets targets def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx], self.targets[idx] dataset MyDataset(...) dataloader DataLoader( dataset, batch_size32, shuffleTrue, num_workers2, )batch_size是每次迭代喂給模型多少條樣本shuffle是在每個 epoch 開始時打亂數(shù)據(jù)順序num_workers是開啟幾個子進程來加載數(shù)據(jù)。需要注意的是在 Windows 上num_workers設置過高有時反而會拖慢速度還容易觸發(fā)奇怪的報錯建議先設成 0 或 1 驗證邏輯再逐步往上調(diào)。3.2 訓練循環(huán)的工程細節(jié)從寫研究 Demo 到寫工程代碼最大的變化是不要把關鍵超參數(shù)寫死在代碼里。學習率、batch size、epoch 數(shù)、模型保存路徑這些東西最好集中放在一個配置里甚至可以直接用命令行參數(shù)控制。這一招在反復調(diào)參、復現(xiàn)實驗時特別管用。一個標準的訓練循環(huán)長這樣model SimpleNet() optimizer torch.optim.Adam(model.parameters(), lr0.001) criterion torch.nn.CrossEntropyLoss() for epoch in range(num_epochs): model.train() total_loss 0.0 for batch_idx, (data, targets) in enumerate(train_loader): data, targets data.to(device), targets.to(device) optimizer.zero_grad() outputs model(data) loss criterion(outputs, targets) loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}/{num_epochs}, Loss: {total_loss / len(train_loader):.4f})這里的model.train()不是擺設它會把模型切換到訓練模式。部分層比如 Dropout 和 BatchNorm在訓練和推理時的行為不一樣你不調(diào)用這個實驗結(jié)果往往會莫名奇妙地變差。同理在推理階段要調(diào)model.eval()并且建議用torch.no_grad()包裹推理過程省顯存還加速。3.3 驗證集不是可選項很多人圖省事直接拿著訓練集上的 loss 判斷模型好壞。這是相當危險的做法。訓練 loss 低只能說明模型記住了這批數(shù)據(jù)并不能說明它對沒見過的數(shù)據(jù)表現(xiàn)如何。正確做法是從訓練集里劃出一部分做驗證集每個 epoch 結(jié)束后在驗證集上評估一次用它作為模型好壞的標尺。劃分方法很簡單比如用torch.utils.data.random_split從數(shù)據(jù)集里隨機取 80% 做訓練20% 做驗證。驗證集的評估代碼基本就是推理代碼注意要加model.eval()和torch.no_grad()。驗證 loss 的價值在于它能暴露“過擬合”的苗頭。當訓練 loss 還在往下降、驗證 loss 開始回升的時候就是模型在死記硬背數(shù)據(jù)而不是學習規(guī)律了。見到這種信號果斷停掉訓練不要跟它死磕。4. 顯存溢出和 loss 不降兩個讓人抓狂的實戰(zhàn)問題4.1 OOM 除了調(diào)小 batch 還能怎么辦顯卡顯存溢出是幾乎每個人都會碰到的報錯提示信息很明確一句 CUDA out of memory 就給你計劃好的下午茶時間取消了。最直接的解決辦法是把batch_size調(diào)小但這不是唯一的思路。我的處理順序通常是這么來的關掉其它占顯存的應用比如瀏覽器里大量標簽頁。檢查代碼里是不是有不該留在顯存的張量用了del之后配合torch.cuda.empty_cache()清一下緩存。調(diào)小batch_size操作系統(tǒng)會自動用更小的內(nèi)存占用計算梯度。如果是訓練大模型可以把optimizer.zero_grad(set_to_noneTrue)配合loss.backward()中間的中間結(jié)果釋放。如果模型本身太大可以考慮用混合精度訓練。PyTorch 官方提供的自動混合精度接口能把刻意不需要高精度的部分用半精度計算顯存占用能明顯下降而且速度往往更快。需要注意的是torch.cuda.empty_cache()只是清空緩存并不是讓占用完全歸零。真正的問題還是要在代碼層面減少駐留在顯存里的張量數(shù)量。4.2 loss 不降甚至變成 NaN遇到 loss 不降不要急著調(diào)參先判斷是哪種不降第一類loss 一直非常大而且不穩(wěn)定甚至出現(xiàn) NaN。這種情況通常是梯度爆炸了。常見誘因是學習率太大、模型初始化不好、或者是 loss 計算里某個值變成無窮大。檢查方法是在訓練循環(huán)里打印梯度范數(shù)total_norm 0.0 for p in model.parameters(): if p.grad is not None: param_norm p.grad.detach().data.norm(2) total_norm param_norm.item() ** 2 total_norm total_norm ** 0.5如果梯度范數(shù)在某個 step 突然飆到非常大基本就鎖定了罪魁禍首。解決辦法是降低學習率或者用梯度裁剪把這些梯度值限制在一定范圍內(nèi)。第二類loss 一直平穩(wěn)降得非常慢或者干脆不動。這有可能是學習率太小也有可能是模型太簡單、表達能力不足?？梢試L試把學習率提高一個數(shù)量級看反應。如果提高了還是沒動靜問題可能出在數(shù)據(jù)預處理上比如沒有做歸一化輸入特征數(shù)值范圍差距太大導致優(yōu)化路徑極其曲折。第三類loss 正常下降但在某個水平就卡住。這種情況往往是模型容量到了瓶頸或者遇到局部最優(yōu)?？梢栽囋囌{(diào)整層數(shù)、寬度或者改用更好的優(yōu)化器比如把樸素 SGD 換成 Adam 或 AdamW。我自己的經(jīng)驗法則是先排除數(shù)據(jù)問題再檢查計算圖邏輯最后才動優(yōu)化器參數(shù)。大部分“玄學”的 loss 問題最后都出在這些基礎環(huán)節(jié)上。4.3 保存報錯時的檢查清單排在訓練中經(jīng)常遇到的問題里還有不少是“從入門到放棄”級別的報錯。比如維度不匹配的錯誤往往是__getitem__返回的數(shù)據(jù)形狀和模型輸入的尺寸對不上。出這種錯別硬猜把每次 batch 的 shape 先在DataLoader里打印出來看一眼。再比如設備不一致的錯誤模型在 GPU 上數(shù)據(jù)還在 CPU 上一算就報錯。我的習慣是在數(shù)據(jù) batch 出來后直接.to(device)并且保證模型也在同一 device 上不要一個地方寫死cuda另一個地方寫cpu。還有一個容易被忽視的問題跨 epoch 的隨機種子。如果你的代碼里用了隨機性操作又沒有固定隨機種子那么每次實驗結(jié)果可能差異都很大。做實驗可以但做復現(xiàn)就麻煩了這一點接下來展開聊。5. 可復現(xiàn)性和斷點續(xù)訓讓實驗最讓人省心的習慣5.1 固定隨機種子深度學習過程里有三處隨機數(shù)據(jù)加載打亂的隨機、模型參數(shù)初始化的隨機、以及某些操作本身比如 Dropout的隨機。為了讓實驗可以對著比有必要在訓練前固定隨機種子import random import numpy as np import torch seed 42 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed)如果有 GPU還想盡量復現(xiàn)可以再加一句torch.cuda.manual_seed_all(seed)需要說明的是不管怎么固定GPU 上某些操作的運算順序仍然可能存在微小的不確定性導致結(jié)果不能 100% 一致。但對于大多數(shù)實驗來說固定種子之后的結(jié)果已經(jīng)具有可比性了不要為了追求絕對復現(xiàn)把自己搞瘋。如果你的數(shù)據(jù)加載用了DataLoader還要確保shuffle的隨機狀態(tài)在每次 epoch 也是可控的通常設置DataLoader的generator參數(shù)就能做到。不過我個人的經(jīng)驗是只要固定了總種子大部分情況下已經(jīng)夠用了除非你要做嚴格對比實驗。5.2 簡單的 checkpoint 保存與恢復訓練 10 個小時結(jié)果中途斷電或者顯存炸一下全部白費這種事我經(jīng)歷過太多次后來一定會寫 checkpoint 邏輯。核心就是把模型狀態(tài)、優(yōu)化器狀態(tài)、當前 epoch、當時的 best loss 都存到一個文件里checkpoint { model_state: model.state_dict(), optimizer_state: optimizer.state_dict(), epoch: epoch, best_loss: best_loss, } torch.save(checkpoint, checkpoint.pt)恢復的時候這樣寫checkpoint torch.load(checkpoint.pt) model.load_state_dict(checkpoint[model_state]) optimizer.load_state_dict(checkpoint[optimizer_state]) start_epoch checkpoint[epoch] 1 best_loss checkpoint[best_loss]這里有幾個關鍵點優(yōu)化器狀態(tài)必須一起保存否則恢復訓練時優(yōu)化器的動量信息會丟失學習率調(diào)度也會對不上相當于從零開始如果你用了學習率調(diào)度器比如StepLR最好也把它的狀態(tài)一起存進去。判斷一個 checkpoint 寫得好不好就看一件事——斷電之后能恢復到幾乎一模一樣的訓練狀態(tài)。我習慣的做法是每訓練完一個 epoch如果驗證 loss 比之前最好記錄更低就覆蓋保存一次。這樣磁盤上永遠只保留最好的模型而不是堆一堆沒用的中間文件。5.3 記錄實驗數(shù)據(jù)的習慣最后說個容易被忽略但又很重要的習慣把每次實驗的關鍵信息記錄下來。我用一個簡單的表格記錄內(nèi)容包括數(shù)據(jù)集取哪一份、模型結(jié)構(gòu)、學習率、batch size、seed、最終驗證精度、跑了多少 epoch、還有什么異?，F(xiàn)象。看上去有點原始但真到了你要回頭找“上次那個效果好的實驗到底用的什么配置”的時候你會發(fā)現(xiàn)比任何花哨的工具都可靠。PyTorch 生態(tài)里有很多現(xiàn)成的實驗管理工具但我始終覺得先把本地手動記錄這件事做好比工具本身更重要。工具可以提高效率卻替代不了對實驗過程的把握。我自己被“這個效果好像復制不出來”坑過很多次最后發(fā)現(xiàn)都是因為當時沒把 seed 和版本號記下來?，F(xiàn)在哪怕只是臨時做一個快速實驗文件夾里也會隨手放一個config.json把該記錄的參數(shù)都寫上。這個小習慣比任何優(yōu)化技巧都更能省時間。