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