搞深度學(xué)習(xí)這幾年身邊但凡問(wèn)“我該學(xué)哪個(gè)框架”的人我?guī)缀醵紩?huì)先反問(wèn)一句你是要做研究、快速驗(yàn)證想法還是直接上生產(chǎn)做高性能推理問(wèn)完這句答案其實(shí)就清晰了一大半。如果以快速迭代、靈活調(diào)試、生態(tài)成熟為主要訴求那pytorch基本是繞不開(kāi)的選擇。今天這篇不打算寫(xiě)成官方文檔式的教程而是從一個(gè)實(shí)際用過(guò)的角度把pytorch的核心機(jī)制、完整訓(xùn)練流程、部署鏈路里那些文檔不會(huì)明說(shuō)、但在實(shí)戰(zhàn)中絕對(duì)會(huì)踩到的點(diǎn)掰開(kāi)揉碎講一遍。這篇文章適合誰(shuí)你可以是完全沒(méi)跑過(guò)深度學(xué)習(xí)項(xiàng)目的新手想搞明白張量、自動(dòng)求導(dǎo)、訓(xùn)練循環(huán)到底是怎么回事也可以是已經(jīng)用其他框架寫(xiě)過(guò)模型、想遷移到pytorch的開(kāi)發(fā)者甚至是在生產(chǎn)環(huán)境里已經(jīng)踩過(guò)一些性能坑、想看看導(dǎo)出和優(yōu)化方案的工程師。我會(huì)盡量做到既有原理層面的解釋也有可以直接抄走的代碼和參數(shù)還會(huì)把我在實(shí)踐中踩過(guò)的坑一并交代清楚。1. 為什么是pytorch從設(shè)計(jì)哲學(xué)到生態(tài)版圖1.1 動(dòng)態(tài)計(jì)算圖讓調(diào)試回歸Python直覺(jué)多數(shù)人第一次接觸pytorch最直觀(guān)的沖擊來(lái)自“動(dòng)態(tài)計(jì)算圖”這個(gè)概念。比喻一下靜態(tài)圖框架像是先畫(huà)好一張完整的電路圖再把電流通進(jìn)去中間每一個(gè)節(jié)點(diǎn)的行為在運(yùn)行前就完全確定了而pytorch的動(dòng)態(tài)圖更像是你在面包板上邊接線(xiàn)邊通電每一根線(xiàn)接完立刻能看到這一小段的輸出。對(duì)于一個(gè)在實(shí)驗(yàn)室里反復(fù)改模型結(jié)構(gòu)的人來(lái)說(shuō)這種“即時(shí)反饋”的價(jià)值幾乎無(wú)法估量。具體到代碼層面動(dòng)態(tài)圖意味著你在前向傳播里寫(xiě)的任何一行Python邏輯——if判斷、for循環(huán)、甚至print調(diào)試——都會(huì)在每次訓(xùn)練迭代時(shí)被真實(shí)執(zhí)行一次。這意味著你可以在前向傳播中間隨意打斷查看某一個(gè)中間張量的shape、數(shù)值分布然后繼續(xù)跑。這種調(diào)試體驗(yàn)和寫(xiě)普通Python程序沒(méi)有區(qū)別也正是很多研究者從其他框架遷移過(guò)來(lái)的核心原因。我記得有一次調(diào)試一個(gè)自定義損失函數(shù)里面有個(gè)索引操作總報(bào)維度錯(cuò)誤我在那個(gè)函數(shù)里隨手加了個(gè)print一眼就看到索引值越界了整個(gè)過(guò)程不到五分鐘。這要是放在靜態(tài)圖環(huán)境里光是定位這個(gè)錯(cuò)誤位置就得折騰半天。當(dāng)然動(dòng)態(tài)圖也有代價(jià)。因?yàn)橛?jì)算圖是逐次構(gòu)建的每次反向傳播的調(diào)度開(kāi)銷(xiāo)相對(duì)更高在某些高性能推理場(chǎng)景下動(dòng)態(tài)圖模式不如靜態(tài)圖優(yōu)化徹底。不過(guò)pytorch早就考慮到了這一點(diǎn)后面用torch.compile和TorchScript做圖優(yōu)化和靜態(tài)化導(dǎo)出這條演進(jìn)路徑已經(jīng)把這些短板補(bǔ)上了后面第4節(jié)里我會(huì)專(zhuān)門(mén)講這部分。1.2 生態(tài)版圖不只是框架是一整套深度學(xué)習(xí)基礎(chǔ)設(shè)施一個(gè)框架能不能長(zhǎng)期用下去看的從來(lái)不只是API好不好寫(xiě)而是生態(tài)是否撐得起從數(shù)據(jù)到部署的整整一條鏈路。pytorch在這點(diǎn)上做得非常成熟我將它常用的部分拆開(kāi)來(lái)看核心庫(kù)與張量運(yùn)算torch本身提供完整的張量操作和自動(dòng)求導(dǎo)這部分是基本功。模型構(gòu)建與訓(xùn)練torch.nn涵蓋常見(jiàn)的層、損失函數(shù)、容器模塊torch.optim有一系列優(yōu)化器torch.utils.data負(fù)責(zé)數(shù)據(jù)加載和預(yù)處理。領(lǐng)域擴(kuò)展庫(kù)TorchVision處理圖像TorchText處理文本TorchAudio處理音頻TorchGeometric處理圖結(jié)構(gòu)數(shù)據(jù)。這些庫(kù)的好處是提供了預(yù)訓(xùn)練權(quán)重和標(biāo)準(zhǔn)數(shù)據(jù)集接口做實(shí)驗(yàn)起步非?？臁２渴鹋c優(yōu)化工具TorchScript把模型靜態(tài)化torch.compile做圖模式編譯加速ONNX導(dǎo)出實(shí)現(xiàn)跨框架互通TorchServe提供模型服務(wù)化方案量化工具鏈支持INT8精度壓縮。這套版圖的一個(gè)核心特點(diǎn)是“從研究到生產(chǎn)平滑過(guò)渡”。你可以在實(shí)驗(yàn)室里用純Python風(fēng)格寫(xiě)一個(gè)科研原型跑通后不用重寫(xiě)整個(gè)邏輯只需要加上導(dǎo)出和優(yōu)化步驟就能部署到C環(huán)境里做低延遲推理。這種連貫性省掉了很多團(tuán)隊(duì)在“研究代碼”和“工程代碼”之間來(lái)回重寫(xiě)的痛苦。2. 核心概念逐個(gè)拆解從張量到訓(xùn)練循環(huán)2.1 張量從NumPy到GPU的一座橋張量Tensor是pytorch里最基礎(chǔ)的數(shù)據(jù)結(jié)構(gòu)它本質(zhì)上是一個(gè)多維數(shù)組。你可以把它理解為“懂得在GPU上并行計(jì)算、并且?guī)е荻刃畔⒌腘umPy數(shù)組”。這里的區(qū)別很關(guān)鍵普通NumPy數(shù)組存放在CPU內(nèi)存里做運(yùn)算時(shí)用多核CPU而pytorch的張量不僅支持CPU還能通過(guò).to(cuda)移動(dòng)到顯存中利用GPU的數(shù)千個(gè)計(jì)算核心做大規(guī)模并行運(yùn)算。實(shí)際寫(xiě)代碼時(shí)最常見(jiàn)的幾個(gè)操作import torch # 從列表或NumPy數(shù)組創(chuàng)建張量 x torch.tensor([[1, 2], [3, 4]]) y torch.from_numpy(np_array) # 從NumPy轉(zhuǎn)換 # 查看張量基本信息 print(x.shape) # torch.Size([2, 2]) print(x.dtype) # torch.int64 print(x.device) # cpu # 移動(dòng)到GPU x_gpu x.to(cuda) # 與NumPy互轉(zhuǎn) x_np x.cpu().numpy()這里有個(gè)特別容易困惑的點(diǎn)torch.tensor()和torch.Tensor()有什么區(qū)別。簡(jiǎn)單說(shuō)torch.tensor()會(huì)根據(jù)傳入數(shù)據(jù)自動(dòng)推斷dtype而torch.Tensor()是默認(rèn)float32的類(lèi)構(gòu)造函數(shù)。新手經(jīng)常在這上面踩坑以為兩者等價(jià)結(jié)果數(shù)據(jù)類(lèi)型不對(duì)后面計(jì)算報(bào)錯(cuò)。個(gè)人習(xí)慣是顯式指定dtype例如torch.tensor(..., dtypetorch.float32)避免歧義。關(guān)于設(shè)備再補(bǔ)充一個(gè)實(shí)操經(jīng)驗(yàn)訓(xùn)練時(shí)不要頻繁做.cpu()和.cuda()之間的切換。每切換一次就涉及一次數(shù)據(jù)拷回拷進(jìn)開(kāi)銷(xiāo)很大有時(shí)候你感覺(jué)訓(xùn)練慢不是算力不夠而是這類(lèi)隱式拷貝太多了。正確的做法是盡量讓整個(gè)計(jì)算鏈都留在同一設(shè)備上只有到打印日志或者保存模型時(shí)才轉(zhuǎn)到CPU。2.2 自動(dòng)求導(dǎo)與計(jì)算圖反向傳播的黑盒是如何被打開(kāi)的深度學(xué)習(xí)訓(xùn)練的核心是反向傳播而pytorch的autograd機(jī)制讓這個(gè)過(guò)程變成了“自動(dòng)發(fā)生”的。當(dāng)你把張量的requires_grad設(shè)為T(mén)rue所有對(duì)這個(gè)張量進(jìn)行的操作都會(huì)被記錄到一張計(jì)算圖中調(diào)用loss.backward()時(shí)pytorch沿著這張圖從輸出往輸入方向計(jì)算每個(gè)葉子節(jié)點(diǎn)的梯度存到各自的.grad里。聽(tīng)起來(lái)有點(diǎn)抽象用代碼演示一下就清楚了x torch.tensor([2.0, 3.0], requires_gradTrue) w torch.tensor([1.0, -1.0], requires_gradTrue) b torch.tensor(0.5, requires_gradTrue) y torch.dot(x, w) b # y 2*1 3*(-1) 0.5 -0.5 loss y ** 2 # loss 0.25 loss.backward() print(x.grad) # tensor([-1., 1.5]) print(w.grad) # tensor([-2., -3.]) print(b.grad) # tensor([-1.])這里每一個(gè)梯度你都可以用手推公式驗(yàn)證。如果對(duì)數(shù)學(xué)敏感你會(huì)發(fā)現(xiàn)它就是在鏈?zhǔn)椒▌t一個(gè)操作一個(gè)操作往回算。整個(gè)過(guò)程中你不需要手動(dòng)實(shí)現(xiàn)任何梯度計(jì)算邏輯這正是深度學(xué)習(xí)框架對(duì)新手最友好的部分。但是自動(dòng)化不等于可以完全不管。有兩個(gè)新手最容易犯的錯(cuò)誤需要特別提醒忘記在每輪迭代前調(diào)用optimizer.zero_grad()。pytorch的設(shè)計(jì)是梯度默認(rèn)累積不清零的話(huà)下一輪迭代會(huì)把上一輪的梯度繼續(xù)累加上去導(dǎo)致參數(shù)更新方向錯(cuò)亂。只要你發(fā)現(xiàn)loss曲線(xiàn)異常波動(dòng)先檢查這里。在不需要求梯度的地方忘了用torch.no_grad()。模型推理和計(jì)算驗(yàn)證指標(biāo)時(shí)根本不需要計(jì)算圖如果還開(kāi)著梯度會(huì)白白占用顯存和算力。推理時(shí)用with torch.no_grad():包裹一下是基本素養(yǎng)。2.3 nn.Module模型組織的世界觀(guān)nn.Module是pytorch組織模型的基類(lèi)。你可以把它看成一個(gè)“容器”里面可以放各種子模塊層、參數(shù)、緩沖區(qū)以及前向傳播邏輯。繼承它之后pytorch會(huì)自動(dòng)追蹤所有參數(shù)的引用這讓model.parameters()、model.to(device)、model.train()/model.eval()這些方法都能開(kāi)箱即用。一個(gè)標(biāo)準(zhǔn)的小模型定義大概是這樣的import torch.nn as nn import torch.nn.functional as F class TwoLayerNet(nn.Module): def __init__(self, in_dim, hidden_dim, out_dim): super().__init__() self.fc1 nn.Linear(in_dim, hidden_dim) self.fc2 nn.Linear(hidden_dim, out_dim) def forward(self, x): x F.relu(self.fc1(x)) x self.fc2(x) return x model TwoLayerNet(784, 256, 10) print(model)我見(jiàn)過(guò)不少新手問(wèn)forward方法里為什么用F.relu而不是nn.ReLU統(tǒng)一回答一下——這兩者功能上幾乎等價(jià)區(qū)別在于F.relu是一個(gè)無(wú)狀態(tài)的函數(shù)調(diào)用適合在forward里即寫(xiě)即用而nn.ReLU是一個(gè)模塊如果你的模型結(jié)構(gòu)里需要顯式地把它作為一個(gè)子層比如帶參數(shù)的模塊鏈那就用nn.ReLU。在不需要保存額外狀態(tài)的情況下用F.relu更簡(jiǎn)潔。在設(shè)計(jì)復(fù)雜模型時(shí)nn.Module的“嵌套”特性非常實(shí)用。你可以把一整塊功能封裝成子模塊再組合進(jìn)主模型。比如一個(gè)Transformer編碼器可以拆成若干層每層又包含自注意力子模塊和前饋?zhàn)幽K層層嵌套調(diào)試時(shí)可以直接針對(duì)某一個(gè)子模塊單獨(dú)打印參數(shù)名和shape定位問(wèn)題快得多。2.4 Dataset與DataLoader把數(shù)據(jù)流做成一條高效管道數(shù)據(jù)加載如果寫(xiě)得不好模型再快也白搭。pytorch提供了torch.utils.data.Dataset和DataLoader這兩個(gè)類(lèi)來(lái)管理數(shù)據(jù)流水線(xiàn)。Dataset負(fù)責(zé)定義“如何讀取一條樣本”DataLoader負(fù)責(zé)把樣本按batch打包并支持多進(jìn)程預(yù)取、隨機(jī)打亂、并行加載。實(shí)現(xiàn)一個(gè)自定義數(shù)據(jù)集需要注意返回值的格式。通常要求__getitem__返回(輸入, 標(biāo)簽)這樣的元組或者字典。如果你的模型輸入不只一個(gè)特征用字典返回會(huì)更清晰比如from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, samples, labels): self.samples samples self.labels labels def __len__(self): return len(self.samples) def __getitem__(self, idx): x torch.tensor(self.samples[idx], dtypetorch.float32) y torch.tensor(self.labels[idx], dtypetorch.long) return x, y loader DataLoader(MyDataset(samples, labels), batch_size32, shuffleTrue, num_workers4)這里num_workers的取值需要跟你機(jī)器的CPU核數(shù)和IO瓶頸配合。并不是越大越好設(shè)得過(guò)高反而會(huì)引起進(jìn)程切換開(kāi)銷(xiāo)和內(nèi)存壓力。我的經(jīng)驗(yàn)法則是先設(shè)成CPU核心數(shù)的一半訓(xùn)練時(shí)觀(guān)察GPU利用率如果利用率上不去再逐步增加。而pin_memoryTrue在GPU訓(xùn)練時(shí)可以加速數(shù)據(jù)從CPU內(nèi)存到顯存的拷貝前提是你確實(shí)在跑CUDA。再提醒一點(diǎn)Dataset.__getitem__里盡量少做昂貴操作。如果你在取樣本時(shí)做圖像解碼、數(shù)據(jù)增強(qiáng)、歸一化這種耗時(shí)工作并且都放在主線(xiàn)程里那么GPU很可能一直在等數(shù)據(jù)。處理圖像數(shù)據(jù)時(shí)更合理的做法是預(yù)處理盡量在離線(xiàn)階段完成在線(xiàn)階段只做輕量級(jí)增強(qiáng)和格式轉(zhuǎn)換。3. 完整實(shí)操流程從零訓(xùn)練一個(gè)圖像分類(lèi)模型3.1 數(shù)據(jù)準(zhǔn)備與增強(qiáng)先讓數(shù)據(jù)“喂”得順下面我用一個(gè)圖像分類(lèi)任務(wù)作為例子帶你完整跑一遍訓(xùn)練流程。數(shù)據(jù)不強(qiáng)求復(fù)雜理解整套邏輯才是重點(diǎn)。第一步自然是準(zhǔn)備數(shù)據(jù)。TorchVision內(nèi)置了若干常用數(shù)據(jù)集直接調(diào)用是很方便的from torchvision import datasets, transforms transform transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) train_dataset datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform)transforms.Normalize里的兩個(gè)參數(shù)分別是均值和標(biāo)準(zhǔn)差它們需要根據(jù)數(shù)據(jù)集統(tǒng)計(jì)得到。對(duì)CIFAR-10這種通用數(shù)據(jù)集官方常見(jiàn)配置就是(0.5, 0.5, 0.5)因?yàn)镽GB三個(gè)通道的像素值歸一化到[0,1]后均值就是0.5。如果你用自己的數(shù)據(jù)集最優(yōu)做法是先統(tǒng)計(jì)全量數(shù)據(jù)的通道均值和標(biāo)準(zhǔn)差再用這些統(tǒng)計(jì)值做歸一化否則歸一化的效果會(huì)打折。數(shù)據(jù)增強(qiáng)方面一個(gè)常用的技巧是“訓(xùn)練集加增強(qiáng)、驗(yàn)證集不加增強(qiáng)”。訓(xùn)練時(shí)做隨機(jī)裁剪、水平翻轉(zhuǎn)可以提升模型泛化能力驗(yàn)證/測(cè)試時(shí)必須保持?jǐn)?shù)據(jù)原始狀態(tài)才能客觀(guān)評(píng)估模型真實(shí)效果。把這一點(diǎn)寫(xiě)清楚在代碼里是很重要的train_transform transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(...), ]) test_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize(...), ])3.2 模型定義與初始化從簡(jiǎn)單結(jié)構(gòu)開(kāi)始定義模型時(shí)我強(qiáng)烈建議你先從一個(gè)小結(jié)構(gòu)開(kāi)始驗(yàn)證整個(gè)流程而不是一上來(lái)就堆復(fù)雜的大網(wǎng)絡(luò)。流程跑通了再換更強(qiáng)的模型排查問(wèn)題的難度會(huì)小很多。下面是一個(gè)簡(jiǎn)單的兩層卷積網(wǎng)絡(luò)class SimpleCNN(nn.Module): def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), nn.Conv2d(32, 64, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), ) self.classifier nn.Sequential( nn.Linear(64 * 8 * 8, 256), nn.ReLU(inplaceTrue), nn.Linear(256, num_classes), ) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) x self.classifier(x) return x這里x.view(x.size(0), -1)是把卷積輸出的多維張量展平成二維矩陣x.size(0)是batch維度-1表示自動(dòng)推斷剩余維度。很多新手在這里算錯(cuò)展平后的維度導(dǎo)致nn.Linear的輸入維度不匹配。其實(shí)你可以先打印x.shape確認(rèn)好維度再寫(xiě)進(jìn)Linear里。關(guān)于參數(shù)初始化pytorch的nn.Linear和nn.Conv2d在創(chuàng)建時(shí)已經(jīng)有默認(rèn)初始化對(duì)大多數(shù)任務(wù)夠用。只有在訓(xùn)練效果不理想、懷疑初始化有問(wèn)題時(shí)才去手動(dòng)改初始化方式比如用nn.init.kaiming_normal_這是針對(duì)ReLU類(lèi)激活函數(shù)的常用策略。3.3 訓(xùn)練循環(huán)一個(gè)標(biāo)準(zhǔn)的迭代流程訓(xùn)練循環(huán)是整篇文章的核心也是新手最容易寫(xiě)錯(cuò)的地方。標(biāo)準(zhǔn)的訓(xùn)練循環(huán)包含四個(gè)步驟前向傳播、計(jì)算損失、反向傳播、更新參數(shù)。順序不能亂尤其不能忘記zero_grad。下面是一個(gè)完整示例import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers2) test_loader DataLoader(test_dataset, batch_size64, shuffleFalse, num_workers2) def run_epoch(model, loader, is_trainTrue): if is_train: model.train() else: model.eval() total_loss, correct, total 0.0, 0, 0 for inputs, labels in loader: inputs, labels inputs.to(device), labels.to(device) if is_train: optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) if is_train: loss.backward() optimizer.step() total_loss loss.item() * inputs.size(0) _, preds torch.max(outputs, dim1) correct (preds labels).sum().item() total inputs.size(0) return total_loss / total, correct / total train_loss, train_acc run_epoch(model, train_loader, is_trainTrue) test_loss, test_acc run_epoch(model, test_loader, is_trainFalse) print(fTrain Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f}) print(fTest Loss: {test_loss:.4f} | Test Acc: {test_acc:.4f})幾個(gè)地方需要解釋一下model.train()與model.eval()的區(qū)別。前者讓Dropout、BatchNorm等層進(jìn)入訓(xùn)練模式后者讓它們進(jìn)入推理模式。只對(duì)BatchNorm和Dropout有影響其他層不受影響但這個(gè)調(diào)用不能省否則你可能會(huì)在驗(yàn)證時(shí)得到一個(gè)異常高的錯(cuò)誤率。loss.item()的用途。loss本身是一個(gè)張量帶計(jì)算圖用它做日志記錄或累加統(tǒng)計(jì)時(shí)會(huì)拖住計(jì)算圖不釋放嚴(yán)重時(shí)會(huì)導(dǎo)致內(nèi)存持續(xù)增長(zhǎng)。item()會(huì)把它轉(zhuǎn)成純Python數(shù)值切斷與計(jì)算圖的關(guān)聯(lián)這樣既省內(nèi)存又方便做float運(yùn)算。為什么要用optimizer.zero_grad()而不是loss.zero_grad()。其實(shí)兩者效果類(lèi)似都可以清空梯度。但標(biāo)準(zhǔn)寫(xiě)法是用optimizer.zero_grad()它會(huì)把所有被優(yōu)化器管理的參數(shù)梯度都清掉邏輯更明確。3.4 驗(yàn)證、保存與加載好的實(shí)驗(yàn)習(xí)慣讓工作不白費(fèi)訓(xùn)練到一定階段后保存模型是一個(gè)重要操作。很多人直接用torch.save(model.state_dict(), model.pth)這是最推薦的方式。它只保存參數(shù)字典不依賴(lài)模型定義的代碼結(jié)構(gòu)加載時(shí)只要你能重新構(gòu)造出同樣的模型類(lèi)再調(diào)用load_state_dict即可# 保存 torch.save(model.state_dict(), model.pth) # 加載 model SimpleCNN() model.load_state_dict(torch.load(model.pth, map_locationcpu)) model.to(device)如果你的模型里定義了自定義類(lèi)加載時(shí)就必須確保這個(gè)類(lèi)在作用域內(nèi)否則會(huì)報(bào)錯(cuò)。這是最常見(jiàn)的一個(gè)保存加載問(wèn)題。所以建議把模型定義和訓(xùn)練腳本分開(kāi)組織或者確保模型類(lèi)所在模塊能被正常import。順便提一下checkpoint策略不只是保存最優(yōu)模型最好每隔若干輪epoch保存一個(gè)包含model.state_dict()、optimizer.state_dict()、epoch、best_acc的完整checkpoint。這樣萬(wàn)一一輪訓(xùn)練崩了你還能從中斷點(diǎn)接著訓(xùn)避免從頭來(lái)過(guò)torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), loss: loss, }, fcheckpoint_epoch{epoch}.pth)4. 從實(shí)驗(yàn)到生產(chǎn)部署鏈路上的關(guān)鍵選擇4.1 模型導(dǎo)出TorchScript與ONNX兩條路線(xiàn)訓(xùn)練完模型只是工作的一半。真正進(jìn)入生產(chǎn)環(huán)境時(shí)你通常會(huì)遇到兩種情況需要在C環(huán)境里做低延遲推理或者需要跟其他框架的推理服務(wù)對(duì)接。這時(shí)就需要把模型從eager模式變成靜態(tài)化表示。TorchScript是將模型靜態(tài)化的官方方案。它有腳本模式和追蹤模式兩種生成方式。追蹤模式用一組示例輸入跑一遍前向把實(shí)際執(zhí)行路徑記錄下來(lái)腳本模式則是用Python語(yǔ)法解析整個(gè)forward邏輯支持控制流。這里有個(gè)重要的注意點(diǎn)如果你的模型里有依賴(lài)數(shù)據(jù)內(nèi)容的if判斷或動(dòng)態(tài)for循環(huán)只用追蹤模式就不夠必須用腳本模式或者把控制流顯式暴露給TorchScript。# 追蹤模式示例 traced_model torch.jit.trace(model, example_input.to(cpu)) traced_model.save(model_traced.pt) # 腳本模式示例用于含控制流的模型 scripted_model torch.jit.script(model) scripted_model.save(model_scripted.pt)ONNX導(dǎo)出則是為了跨框架互操作性。把pytorch模型導(dǎo)出為ONNX后可以轉(zhuǎn)到ONNX Runtime、TensorRT、OpenVINO等推理引擎也能被一些云端推理平臺(tái)直接使用dummy_input torch.randn(1, 3, 32, 32) torch.onnx.export( model.cpu(), dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, )dynamic_axes可以讓導(dǎo)出的模型支持可變batch維度非常適合在線(xiàn)推理場(chǎng)景。如果不設(shè)置導(dǎo)出的模型會(huì)固定為batch1換成其他batch尺寸就得重新導(dǎo)出。4.2 性能優(yōu)化torch.compile與量化對(duì)于訓(xùn)練階段的加速torch.compile是現(xiàn)代pytorch的大殺器。它通過(guò)在首次運(yùn)行時(shí)對(duì)模型做圖級(jí)優(yōu)化把多個(gè)kernel融合起來(lái)減少GPU顯存讀寫(xiě)和調(diào)度開(kāi)銷(xiāo)。使用方式極其簡(jiǎn)單一行代碼model torch.compile(model)注意torch.compile在首次調(diào)用時(shí)有編譯開(kāi)銷(xiāo)后續(xù)才會(huì)快起來(lái)。而且它跟某些自定義算子、動(dòng)態(tài)控制流可能存在兼容性問(wèn)題。實(shí)踐中如果遇到奇怪的錯(cuò)誤可以先試試關(guān)掉編譯再排查是不是它引起的。推理階段的性能優(yōu)化量化是更常用的手段。將模型從FP32降到INT8顯存占用和推理延遲通常都能顯著下降。常見(jiàn)做法是訓(xùn)練后量化quantized_model torch.quantization.quantize_dynamic( model.cpu(), {nn.Linear, nn.Conv2d}, dtypetorch.qint8 )quantize_dynamic只量化權(quán)重在推理時(shí)動(dòng)態(tài)計(jì)算量化后的激活值實(shí)現(xiàn)簡(jiǎn)單、精度損失小適合快速落地。如果要追求更高壓縮率或更低延遲需要做全量化也就是帶校準(zhǔn)數(shù)據(jù)的靜態(tài)量化那就要額外準(zhǔn)備一個(gè)校準(zhǔn)數(shù)據(jù)集流程復(fù)雜一些收益也更明顯。4.3 推理服務(wù)與落地陷阱生產(chǎn)環(huán)境部署模型除了模型本身還要考慮批處理、動(dòng)態(tài)shape、并發(fā)保護(hù)、超時(shí)等工程問(wèn)題。我見(jiàn)過(guò)太多團(tuán)隊(duì)花了巨大精力優(yōu)化模型網(wǎng)絡(luò)結(jié)構(gòu)結(jié)果被一個(gè)低效的推理封裝拖累了整體性能。幾個(gè)實(shí)用的建議盡量用批量推理而不是單條推理。GPU的算力特性決定了它更擅長(zhǎng)處理大批量。即使在線(xiàn)服務(wù)一次只來(lái)一個(gè)請(qǐng)求也可以做成攢一批再推理的機(jī)制。如果請(qǐng)求量不大批量帶來(lái)的延遲提升通常比你想的大得多。對(duì)固定shape做優(yōu)化。很多推理引擎對(duì)靜態(tài)shape的優(yōu)化比對(duì)動(dòng)態(tài)shape好得多。如果業(yè)務(wù)場(chǎng)景允許盡量把輸入尺寸固定。如果你用ONNX導(dǎo)出時(shí)設(shè)置了dynamic_axes在某些推理引擎上會(huì)性能下降這時(shí)就需要權(quán)衡靈活性還是性能。不要忽略CPU推理場(chǎng)景。不是所有部署環(huán)境都有GPU。在CPU上pytorch的推理性能不如專(zhuān)門(mén)優(yōu)化過(guò)的ONNX Runtime或OpenVINO這也是導(dǎo)出ONNX模型的一個(gè)重要價(jià)值——同樣的模型在正確的推理引擎上能跑出完全不同的性能。5. 高頻問(wèn)題與排查實(shí)戰(zhàn)5.1 張量形狀不匹配最常見(jiàn)的報(bào)錯(cuò)RuntimeError: size mismatch這類(lèi)錯(cuò)誤本質(zhì)是模型的前向傳播中某個(gè)操作要求張量形狀匹配但實(shí)際傳進(jìn)來(lái)的不匹配。排查思路其實(shí)很簡(jiǎn)單就一條定位是哪一行報(bào)錯(cuò)然后打印這一步的輸入和輸出shape。# 在可疑操作之前打印 print(before fc1:, x.shape) x self.fc1(x) print(after fc1:, x.shape)很多新手一看到shape不匹配就慌了其實(shí)這類(lèi)問(wèn)題90%都出在nn.Linear的輸入維度和上一層輸出維度不一致上。經(jīng)驗(yàn)法則是從數(shù)據(jù)進(jìn)入模型的第一個(gè)操作開(kāi)始逐步打印shape直到找到第一個(gè)不匹配的地方。不要猜打印出來(lái)比什么都直接。5.2 顯存溢出OOM排查清單CUDA out of memory是訓(xùn)練大模型時(shí)的噩夢(mèng)。遇到這個(gè)報(bào)錯(cuò)按以下順序排查看代碼里是否有顯存泄漏。最容易出現(xiàn)的地方是驗(yàn)證階段的loss.item()用得不對(duì)或者推理時(shí)忘了torch.no_grad()導(dǎo)致每輪驗(yàn)證都往顯存里塞一堆計(jì)算圖。這種情況一旦發(fā)生顯存會(huì)隨著迭代輪數(shù)持續(xù)增長(zhǎng)直到爆掉。降低batch size。這是最直接的臨時(shí)緩解方式。batch size減半顯存占用大約是原來(lái)的一半多一點(diǎn)因?yàn)檫€有模型參數(shù)和優(yōu)化器狀態(tài)的固定開(kāi)銷(xiāo)。減少后如果訓(xùn)練不穩(wěn)定可以相應(yīng)調(diào)低學(xué)習(xí)率。檢查梯度累積。如果batch size確實(shí)不能再小就用梯度累積模擬更大batch。這里的關(guān)鍵點(diǎn)又回到zero_grad累積時(shí)不能每一小步都step而是多走幾步后再更新參數(shù)注意最后一定要除累積步數(shù)。清理中間激活。如果模型很深可以在不需要中間結(jié)果的層之間手動(dòng)del一些大張量或者用torch.utils.checkpoint對(duì)部分層做激活檢查點(diǎn)用計(jì)算換內(nèi)存。5.3 數(shù)據(jù)加載慢與DataLoader調(diào)優(yōu)訓(xùn)練時(shí)GPU利用率很低但顯存沒(méi)有爆大概率是數(shù)據(jù)加載拖后腿。你可以做一個(gè)簡(jiǎn)單測(cè)試在循環(huán)里只加載數(shù)據(jù)不跑前向和反向看每秒能取多少batch。如果瓶頸在數(shù)據(jù)側(cè)可以考慮下面幾個(gè)措施增加num_workers。多進(jìn)程并行讀取通常能顯著提升吞吐率。開(kāi)啟pin_memoryTrue。在GPU訓(xùn)練時(shí)它可以減少CPU到GPU的數(shù)據(jù)拷貝時(shí)間。檢查數(shù)據(jù)預(yù)處理是否太重。如果__getitem__里有昂貴的圖像解碼或數(shù)據(jù)增強(qiáng)考慮把一部分預(yù)處理遷移到離線(xiàn)階段完成或者用更高效的庫(kù)例如用pillow配合opencv處理圖像時(shí)優(yōu)先用支持底層SIMD加速的實(shí)現(xiàn)。5.4 梯度相關(guān)坑zero_grad、梯度累積與梯度裁剪梯度問(wèn)題在調(diào)試訓(xùn)練時(shí)特別常見(jiàn)。除了忘記zero_grad導(dǎo)致loss發(fā)散還有幾個(gè)隱蔽場(chǎng)景梯度累積時(shí)忘了歸一化。用梯度累積模擬大batch時(shí)如果累積了4個(gè)小batch最終更新時(shí)的梯度相當(dāng)于4個(gè)batch的和所以要把總梯度除以4否則學(xué)習(xí)率等效于被放大了4倍模型容易震蕩。梯度爆炸。在訓(xùn)練RNN或Transformer這類(lèi)模型時(shí)梯度范數(shù)可能突然暴漲導(dǎo)致loss變成NaN。常規(guī)做法是梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)max_norm的經(jīng)驗(yàn)值在0.5到2.0之間。設(shè)太大了等于沒(méi)設(shè)設(shè)太小了會(huì)拖慢收斂需要根據(jù)實(shí)際觀(guān)察調(diào)整。為什么有時(shí)候requires_gradTrue也有g(shù)rad為None如果一個(gè)葉子節(jié)點(diǎn)不參與計(jì)算圖中對(duì)loss的任何路徑它的grad自然就是None。檢查一下你的模型里是否有參數(shù)被意外凍住或者該參數(shù)所在的子模塊沒(méi)被調(diào)用即可。6. 最后一個(gè)想說(shuō)的實(shí)踐技巧就在我準(zhǔn)備收尾的時(shí)候突然想到一個(gè)經(jīng)常被忽略、卻能實(shí)實(shí)在在省時(shí)間的點(diǎn)把實(shí)驗(yàn)配置全部參數(shù)化。我見(jiàn)過(guò)太多人在腳本里硬編碼學(xué)習(xí)率、batch size、epoch數(shù)每換一組參數(shù)就改一遍代碼既不安全也不方便追蹤。解決方式很簡(jiǎn)單用一個(gè)字典或者簡(jiǎn)潔的數(shù)據(jù)類(lèi)把超參集中管理config { lr: 1e-3, batch_size: 64, epochs: 20, seed: 42, model: SimpleCNN, }固定隨機(jī)種子這一點(diǎn)尤其值得說(shuō)。深度學(xué)習(xí)訓(xùn)練有大量隨機(jī)性來(lái)源——數(shù)據(jù)加載順序、初始化、Dropout等。如果不固定種子你跑兩次實(shí)驗(yàn)得到的結(jié)果可能差異不小這會(huì)讓你很難判斷某個(gè)改動(dòng)到底有沒(méi)有效果。在腳本開(kāi)頭加上import random import numpy as np import torch def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)這算是我踩過(guò)多次坑后養(yǎng)成的習(xí)慣。有時(shí)候你以為模型結(jié)構(gòu)改對(duì)了結(jié)果只是運(yùn)氣差換個(gè)種子效果就回去了。固定種子之后至少能保證實(shí)驗(yàn)對(duì)比的公平性?；氐介_(kāi)頭那個(gè)問(wèn)題做研究、快速驗(yàn)證、靈活調(diào)試pytorch確實(shí)是我當(dāng)前最推薦的選擇。它的動(dòng)態(tài)計(jì)算圖、健全的生態(tài)和從實(shí)驗(yàn)到部署的平滑過(guò)渡讓深度學(xué)習(xí)項(xiàng)目的整個(gè)生命周期都能在一個(gè)框架里跑通。希望這篇從原理到實(shí)操再到部署和排查的梳理能讓你少走一些我走過(guò)的彎路。