戰(zhàn)：從數(shù)據(jù)處理到遷移學(xué)習(xí)的完整流程)
1. 項目概述從一張圖片到識別結(jié)果到底發(fā)生了什么用 PyTorch 做圖片識別核心任務(wù)通常是圖像分類也就是我們輸入一張圖片模型輸出它屬于哪個類別。比如一張貓的照片經(jīng)過模型前向傳播得到貓 0.92、狗 0.05、鳥 0.02這樣的概率分布最后取最高概率對應(yīng)的類別作為預(yù)測結(jié)果。很多初學(xué)者拿到這個需求時第一反應(yīng)是調(diào)個模型跑一下但實(shí)際落地時你會發(fā)現(xiàn)整個流程是一條完整鏈路圖像讀取、預(yù)處理、數(shù)據(jù)加載、模型搭建、訓(xùn)練迭代、驗證調(diào)參、固定權(quán)重、部署推理。任何一個環(huán)節(jié)處理不當(dāng)訓(xùn)練出來的模型精度都會差得離譜。這篇文章我會把這條鏈路每一個環(huán)節(jié)的實(shí)操細(xì)節(jié)都拆開講清楚包括一些你翻文檔翻不到的經(jīng)驗和坑。這篇文章適合誰看剛接觸 PyTorch、想快速跑通一個圖像分類項目的同學(xué)或者已經(jīng)能跑通官方 Demo 但換了自己的數(shù)據(jù)集就不知所措的朋友。我會用貓狗分類作為貫穿全文的示例任務(wù)因為它的數(shù)據(jù)獲取簡單、類別直觀、訓(xùn)練收益明顯。文章里涉及的概念和代碼寫法換成任何其他分類任務(wù)都通用。2. 核心鏈路拆解圖片識別的五段流程2.1 圖像分類的本質(zhì)圖像分類是計算機(jī)視覺最基礎(chǔ)的任務(wù)。它的數(shù)學(xué)本質(zhì)是學(xué)習(xí)一個映射函數(shù) f輸入一張 H×W×C 的圖片張量輸出一個長度為 K 的向量K 是類別總數(shù)向量中的值經(jīng)過 Softmax 歸一化后表示每個類別的概率。在這里有一個很多新手沒想明白的點(diǎn)圖片輸入到網(wǎng)絡(luò)里其實(shí)是一堆數(shù)字。比如一張 224×224 的 RGB 圖片輸入張量的形狀是 [3, 224, 224]取值范圍在 0 到 1 之間。卷積神經(jīng)網(wǎng)絡(luò)在做的事情本質(zhì)上是從這些數(shù)字中逐層提取特征——淺層卷積提取邊緣、顏色、紋理深層卷積提取眼睛、耳朵、輪子這種語義級別的部件。最后一層全連接將這些特征映射到類別空間。所以整個識別流程的五個環(huán)節(jié)是圖像預(yù)處理 → 數(shù)據(jù)加載 → 模型前向傳播 → 損失計算與反向傳播 → 推理預(yù)測。環(huán)境準(zhǔn)備完成后后面每個環(huán)節(jié)我會單獨(dú)展開講。2.2 環(huán)境準(zhǔn)備與硬件選型先說你需要的環(huán)境這里給出一個我用下來比較穩(wěn)定的版本組合PyTorch 2.x CUDA 11.8 或 12.1 Python 3.9~3.11 torchvision 配套版本。嚴(yán)格來說我們不建議追求最新版本穩(wěn)定性和兼容性更重要很多 GPU 驅(qū)動對最新的 CUDA 版本支持滯后跑起來反而出各種詭異報錯。有兩點(diǎn)經(jīng)驗說明一下。第一如果你用的是 NVIDIA 顯卡安裝 PyTorch 之前先去官網(wǎng)根據(jù)你的 CUDA 版本選擇對應(yīng)的安裝命令不要直接用 pip install torch 默認(rèn)裝 CPU 版本不然訓(xùn)練速度會慢幾十倍。第二如果沒有獨(dú)立顯卡也可以先跑通整個流程只是訓(xùn)練時間會長很多本文所有代碼在 CPU 上都能運(yùn)行只不過我們建議把 batch size 調(diào)小。3. 數(shù)據(jù)處理與加載比模型更影響精度的地方3.1 數(shù)據(jù)目錄如何組織PyTorch 的 torchvision.datasets.ImageFolder 是處理分類任務(wù)最省事的工具前提是你的數(shù)據(jù)目錄按類別組織成這樣的結(jié)構(gòu)data/ train/ cat/ cat_001.jpg cat_002.jpg dog/ dog_001.jpg dog_002.jpg val/ cat/ cat_001.jpg dog/ dog_001.jpgtrain 和 test 文件夾各自包含類別子文件夾里面的圖片按類別放置。用這種方式組織數(shù)據(jù)ImageFolder 會自動讀取所有圖片并根據(jù)子文件夾名稱生成類別索引我們完全可以用自己真實(shí)的數(shù)據(jù)目錄來替換這個示例結(jié)構(gòu)。一個要注意的點(diǎn)訓(xùn)練集和驗證集必須分開。我見過不少朋友把所有圖片放在一個大文件夾里直接訓(xùn)練這樣驗證集失去了意義。因為模型見過了所有數(shù)據(jù)驗證集精度虛高換到真實(shí)場景性能崩盤。劃分比例一般是 82 或 91。3.2 Dataset 與 DataLoader 的關(guān)鍵配置有了目錄結(jié)構(gòu)接下來就是標(biāo)準(zhǔn)的數(shù)據(jù)加載代碼。這段代碼幾乎是所有 PyTorch 圖像分類項目的骨架你需要根據(jù)自己的數(shù)據(jù)集特點(diǎn)去調(diào)整里面的參數(shù)。import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # 訓(xùn)練集預(yù)處理 train_transform transforms.Compose([ transforms.Resize((256, 256)), # 先縮放 transforms.RandomCrop(224), # 再隨機(jī)裁剪 transforms.RandomHorizontalFlip(), # 隨機(jī)水平翻轉(zhuǎn) transforms.ToTensor(), # 轉(zhuǎn)為張量并歸一化到 [0,1] transforms.Normalize( mean[0.485, 0.456, 0.406], # ImageNet 均值 std[0.229, 0.224, 0.225] # ImageNet 標(biāo)準(zhǔn)差 ) ]) # 驗證集預(yù)處理只做縮放和歸一化不做數(shù)據(jù)增強(qiáng) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) train_dataset datasets.ImageFolder(data/train, transformtrain_transform) val_dataset datasets.ImageFolder(data/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue)這里面的重點(diǎn)在于為什么訓(xùn)練要用縮放 隨機(jī)裁剪而驗證集直接縮放。訓(xùn)練集需要數(shù)據(jù)增強(qiáng)來增加樣本多樣性隨機(jī)裁剪相當(dāng)于在 256×256 的圖片中隨機(jī)截取 224×224 的區(qū)域相當(dāng)于一個樣本生成了很多不同的子樣本能有效緩解過擬合。而驗證集的目的是評估模型的真實(shí)表現(xiàn)所以只用固定縮放保證一致性。另一個容易忽略的配置是 shuffle 參數(shù)。訓(xùn)練集必須 shuffleTrue因為如果不打亂每個 batch 里都是同一類別的圖片梯度下降的方向會來回震蕩收斂速度慢且不穩(wěn)定。驗證集則不需要 shuffle。關(guān)于 Normalize 的 mean 和 std這一組值來自 ImageNet 數(shù)據(jù)集的統(tǒng)計結(jié)果。我們做遷移學(xué)習(xí)時通常沿用這組值因為它和預(yù)訓(xùn)練模型的輸入分布是匹配的。如果從頭訓(xùn)練自己的模型理論上應(yīng)該重新統(tǒng)計自己數(shù)據(jù)集的均值和標(biāo)準(zhǔn)差但實(shí)際操作中大多數(shù)情況下用了 ImageNet 的數(shù)值也不會差太多。3.3 數(shù)據(jù)增強(qiáng)做多少才算合適數(shù)據(jù)增強(qiáng)是把雙刃劍。太少模型容易過擬合太多會給訓(xùn)練制造不必要的難度。以貓狗分類為例我常用的增強(qiáng)組合按優(yōu)先級排列隨機(jī)水平翻轉(zhuǎn)必選不影響語義隨機(jī)裁剪必選模擬尺度變化隨機(jī)顏色抖動可選增強(qiáng)光照魯棒性隨機(jī)旋轉(zhuǎn)最多 ±15 度太大會破壞圖像結(jié)構(gòu)如果數(shù)據(jù)集只有幾千張圖片我建議增強(qiáng)力度適當(dāng)加強(qiáng)把 RandomResizedCrop 的尺度范圍設(shè)為 (0.5, 1.0)讓模型學(xué)會從不同尺度下識別目標(biāo)。如果數(shù)據(jù)集有十萬級別增強(qiáng)可以適度保守重點(diǎn)放在防止過擬合上。我做項目時驗證過一個經(jīng)驗當(dāng)訓(xùn)練集樣本量少于 5000 時數(shù)據(jù)增強(qiáng)的收益是顯著的分類精度可以提升 2 到 5 個百分點(diǎn)。但如果你用的樣本量已經(jīng)很大或者數(shù)據(jù)本身質(zhì)量很高盲目加增強(qiáng)反而可能導(dǎo)致精度不升反降。4. 模型如何選手寫 CNN 還是遷移學(xué)習(xí)4.1 從零搭建一個可用的 CNN如果從頭搭建一個卷積神經(jīng)網(wǎng)絡(luò)核心是感受野的設(shè)計。簡單來說每一層卷積核覆蓋的區(qū)域大小決定了網(wǎng)絡(luò)能看到多大范圍的圖像信息。一個經(jīng)典的小型 CNN 結(jié)構(gòu)如下import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x這個模型的思路是輸入 3×224×224 圖片經(jīng)過三層卷積 池化后特征圖變成 128×28×28然后通過 AdaptiveAvgPool2d 把特征圖壓縮成 1×1再接全連接分類。三層卷積提取不同層次的特征池化逐步降低分辨率擴(kuò)大感受野。從零搭建 CNN 的優(yōu)勢是不依賴外部權(quán)重部署時自由度很大訓(xùn)練過程完全可控。缺點(diǎn)是精度天花板明顯對于復(fù)雜任務(wù)很難超過預(yù)訓(xùn)練模型。所以我的建議是如果數(shù)據(jù)集很小但圖像特征明顯比如純色背景下區(qū)分圓形和方形自建 CNN 足夠用如果圖像復(fù)雜、類別相近優(yōu)先考慮遷移學(xué)習(xí)。4.2 遷移學(xué)習(xí)用預(yù)訓(xùn)練模型大幅提升精度遷移學(xué)習(xí)的邏輯和自建模型完全不同。它的核心是復(fù)用在大規(guī)模數(shù)據(jù)集上訓(xùn)練好的模型參數(shù)——比如在 ImageNet 上已經(jīng)學(xué)好了邊緣、紋理、形狀等通用特征——然后在自己數(shù)據(jù)集上微調(diào)。這就像你不需要重新學(xué)習(xí)認(rèn)識顏色和輪廓只需要學(xué)習(xí)貓和狗的區(qū)別就行。使用 torchvision 的預(yù)訓(xùn)練模型最常用的幾個選項ResNet18 / ResNet50精度和速度最均衡的選擇EfficientNet-B0計算量更小精度相當(dāng)MobileNetV3適合移動端部署場景VGG16結(jié)構(gòu)簡單但參數(shù)量大訓(xùn)練慢實(shí)用性下降以 ResNet18 為例遷移學(xué)習(xí)的標(biāo)準(zhǔn)寫法是凍結(jié)特征層參數(shù)替換最后一層全連接from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 凍結(jié)所有層的參數(shù) for param in model.parameters(): param.requires_grad False # 獲取最后一層全連接的輸入維度 num_features model.fc.in_features # 替換成自己的分類器 model.fc nn.Sequential( nn.Linear(num_features, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, 2) )這段代碼里最關(guān)鍵的邏輯是 requires_grad False。凍結(jié)特征層參數(shù)后反向傳播的梯度不會更新這些層的權(quán)重模型只會訓(xùn)練我們新加的全連接層。這樣有幾個好處一是訓(xùn)練速度大幅提升因為不需要計算大量參數(shù)的梯度二是可以有效防止數(shù)據(jù)集過小導(dǎo)致的過擬合避免破壞預(yù)訓(xùn)練模型已經(jīng)學(xué)好的通用特征表示。微調(diào)有兩種常用策略一是上面這種只訓(xùn)練分類頭的方案適合數(shù)據(jù)集幾百到幾千張的情況二是全量微調(diào)也就是所有層都放開訓(xùn)練適合數(shù)據(jù)集規(guī)模較大時有足夠的樣本來刷新預(yù)訓(xùn)練知識。實(shí)際項目中我建議先用凍結(jié)方式訓(xùn)練幾個 epoch 讓分類頭收斂然后解凍最后幾層做微調(diào)這樣精度更高也穩(wěn)定。4.3 損失函數(shù)和優(yōu)化器的選擇邏輯分類任務(wù)的標(biāo)準(zhǔn)配置是交叉熵?fù)p失函數(shù)加 SGD 或 Adam。這里的交叉熵?fù)p失函數(shù)值得多說一句PyTorch 的 nn.CrossEntropyLoss 其實(shí)做了兩件事——對模型的輸出做 Softmax 歸一化然后計算負(fù)對數(shù)似然損失。所以模型最后一層不需要再接 Softmax直接輸出原始 logits 就好否則等于做了兩次歸一化精度會受影響。criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.fc.parameters(), lr0.001)優(yōu)化器方面Adam 適合快速迭代和調(diào)參對學(xué)習(xí)率的敏感度低一些。SGD 加 Momentum 需要更精細(xì)的學(xué)習(xí)率調(diào)整但收斂后的精度往往更高。我在實(shí)際訓(xùn)練中習(xí)慣先用 Adam 跑通流程之后如果有優(yōu)化需求再切換到帶學(xué)習(xí)率衰減的 SGD。這里有個新手常踩的坑如果只訓(xùn)練分類頭優(yōu)化的參數(shù)只有 model.fc.parameters()因為被凍結(jié)的層沒有梯度。如果全量微調(diào)則要傳入 model.parameters()。如果混用凍結(jié)層和可訓(xùn)練層需要確保優(yōu)化器中傳入了正確的參數(shù)組否則會出現(xiàn)模型一點(diǎn)也不學(xué)習(xí)的困惑現(xiàn)象。4.4 設(shè)備準(zhǔn)備讓代碼自動跑在 GPU 或 CPU 上模型訓(xùn)練前要先確定計算設(shè)備。建議直接寫成下面這種自適應(yīng)方式避免在 GPU 機(jī)器上調(diào)試過的代碼挪到 CPU 機(jī)器上直接報錯device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)model.to(device) 這一步很多人會漏掉導(dǎo)致前向傳播時模型還在 CPU 上數(shù)據(jù)在 GPU 上報錯說 device mismatch。還有一點(diǎn)需要特別注意輸入圖片張量也要調(diào)用 .to(device) 移動設(shè)備我在初學(xué)階段經(jīng)常因為忘記把 data 移動到 GPU 上而報錯這種問題排查起來很浪費(fèi)時間。5. 訓(xùn)練循環(huán)與模型驗證的完整實(shí)現(xiàn)5.1 標(biāo)準(zhǔn)訓(xùn)練循環(huán)代碼訓(xùn)練循環(huán)是整個項目最核心的代碼段。我把一個完整的訓(xùn)練驗證循環(huán)貼出來這份代碼可以直接用以下方式運(yùn)行然后根據(jù)打印的日志判斷訓(xùn)練效果。def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc def validate(model, val_loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc一個容易犯錯的地方是寫法上訓(xùn)練時必須調(diào)用 model.train()驗證時必須調(diào)用 model.eval()。這兩個模式的區(qū)別在于 BatchNorm 和 Dropout 的行為train 模式下 Dropout 會隨機(jī)丟棄神經(jīng)元BatchNorm 會使用 batch 內(nèi)的統(tǒng)計量eval 模式下 Dropout 關(guān)閉BatchNorm 使用訓(xùn)練階段累積的全局統(tǒng)計量。如果驗證時忘了切換到 eval 模式驗證精度會異常波動并且吃進(jìn)隨機(jī)性。另外驗證階段必須包裹 torch.no_grad()它有兩個用處一是停止梯度追蹤節(jié)省顯存二是避免驗證過程意外修改模型參數(shù)。5.2 模型保存策略只保留最優(yōu)權(quán)重訓(xùn)練過程中的模型保存也是有講究的。不經(jīng)思考地只在最后保存一次模型如果你中途過擬合了最后保存的就是過擬合的模型。我建議采用驗證精度最高時保存的策略best_acc 0.0 for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(...) print(fEpoch {epoch1}: train_loss{train_loss:.4f}, train_acc{train_acc:.4f}, fval_loss{val_loss:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(f - 保存最優(yōu)模型驗證精度 {best_acc:.4f})torch.save 建議只保存 model.state_dict()它只保存模型權(quán)重參數(shù)不攜帶完整的模型結(jié)構(gòu)。原因有兩個一是文件體積更小二是加載時更靈活。你在 state_dict 之外還需要單獨(dú)記錄類別映射關(guān)系否則推理時只有數(shù)字類別沒有類別名稱。這個我后面會講。5.3 學(xué)習(xí)率調(diào)整和 early stopping 的實(shí)際價值訓(xùn)練過程中學(xué)習(xí)率設(shè)置是一個常見痛點(diǎn)。學(xué)習(xí)率太大模型震蕩不收斂太小則訓(xùn)練半天損失紋絲不動。我喜歡用 ReduceLROnPlateau 做驗證精度的自適應(yīng)調(diào)整scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience3 ) for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(...) scheduler.step(val_loss)這里傳入的 modemin 表示當(dāng) val_loss 不再下降時把學(xué)習(xí)率乘以 0.5。patience3 意味著連續(xù) 3 個 epoch 驗證損失沒有明顯下降才觸發(fā)衰減。學(xué)習(xí)率衰減的原理是訓(xùn)練初期需要較大的步長快速靠近最優(yōu)點(diǎn)但接近最優(yōu)點(diǎn)后大步長會在最優(yōu)值附近震蕩此時縮小步長可以讓損失繼續(xù)下降。early stopping 指的是當(dāng)驗證損失連續(xù)多個 epoch 不再下降時提前終止訓(xùn)練。這個和保存最優(yōu)模型配合起來用能有效避免過擬合。但我實(shí)際使用中有一個心得early stopping 的 patience 設(shè)置要根據(jù)數(shù)據(jù)集大小靈活調(diào)整小數(shù)據(jù)集設(shè) 5 個 epoch 比較合理大數(shù)據(jù)集可能需要 10 個甚至更多 epoch 的耐心否則容易在驗證損失的局部極小值處提前停止。6. 推理階段加載模型識別單張新圖片6.1 完整的推理代碼訓(xùn)練完成后真正到識別圖片這個動作時代碼要包含三個步驟加載模型權(quán)重、對輸入圖片做和訓(xùn)練時完全一致的預(yù)處理、前向傳播得到預(yù)測結(jié)果。這份推理代碼我是按照生產(chǎn)環(huán)境的標(biāo)準(zhǔn)寫的可以直接用import torch from PIL import Image from torchvision import transforms def predict_image(image_path, model, device, class_names): # 1. 加載并預(yù)處理圖片 image Image.open(image_path).convert(RGB) # 2. 預(yù)處理與驗證集的預(yù)處理完全一致 preprocessing transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) input_tensor preprocessing(image).unsqueeze(0) # 形狀 [1, 3, 224, 224] input_tensor input_tensor.to(device) # 3. 推理 model.eval() with torch.no_grad(): outputs model(input_tensor) probabilities torch.softmax(outputs, dim1) confidence, predicted_idx torch.max(probabilities, dim1) predicted_class class_names[predicted_idx.item()] return predicted_class, confidence.item()調(diào)用示例# 設(shè)定類別映射0 對應(yīng) cat1 對應(yīng) dog class_names [cat, dog] # 加載訓(xùn)練好的最優(yōu)權(quán)重 model models.resnet18(weightsNone) num_features model.fc.in_features model.fc nn.Sequential(...) # 結(jié)構(gòu)必須和訓(xùn)練時完全一致 model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model model.to(device) result, confidence predict_image(test_photo.jpg, model, device, class_names) print(f識別結(jié)果{result}置信度{confidence:.2%})這里的關(guān)鍵坑在于 load_state_dict 前模型結(jié)構(gòu)必須與訓(xùn)練時完全一致。如果你在訓(xùn)練代碼里改過 model.fc 的結(jié)構(gòu)那么推理代碼里也要做一模一樣的修改否則加載權(quán)重時會報 size mismatch。為了避免這種問題我強(qiáng)烈建議把模型構(gòu)建代碼寫在一個函數(shù)里訓(xùn)練和推理都調(diào)用同一個函數(shù)不要二次復(fù)制。6.2 置信度識別結(jié)果的第二落點(diǎn)很多人做圖片識別只看預(yù)測類別但置信度信息其實(shí)更具價值。置信度是 Softmax 輸出的最大概率值表示為模型對該預(yù)測的自信程度。我在實(shí)際項目中總結(jié)的置信度決策規(guī)則置信度 0.9模型非常確定可以放心采納置信度在 0.6 到 0.9 之間模型有一定把握可以采納但要提示用戶核實(shí)置信度低于 0.6建議返回?zé)o法確定因為多類別中每個類別的概率都比較分散強(qiáng)取最大值很可能是錯誤判斷一個常見場景是模型對貓的置信度是 0.45對狗的置信度是 0.42此時取貓作為結(jié)果并不穩(wěn)妥。這種低置信度通常意味著圖片里的目標(biāo)不明顯、遮擋嚴(yán)重或包含訓(xùn)練數(shù)據(jù)中未出現(xiàn)的背景干擾。如果業(yè)務(wù)場景對誤判容忍度低建議加一個置信度閾值判斷。6.3 推理性能優(yōu)化經(jīng)驗如果識別功能要集成到實(shí)際應(yīng)用中推理速度和延遲就很重要了。幾個實(shí)用的優(yōu)化方向把整個模型用 model.eval() 切換去掉 Dropout 的隨機(jī)性對輸入圖片做 batch 推理把多張圖片合成一個 batch 一次前向傳播比單張循環(huán)快很多倍使用 torch.no_grad() 防止構(gòu)建計算圖如果需要更高吞吐量可以考慮用 ONNX 導(dǎo)出模型加速但在精度要求高時要注意導(dǎo)出和推理的精度差異我參與過的一個圖片識別的技術(shù) Demo用 ResNet18 在 GPU 上單張推理延遲只有幾毫秒完全滿足實(shí)時應(yīng)用的需求。而如果換用 CPU 推理延遲會上升到幾十毫秒這種差距會讓一些對時延敏感的應(yīng)用場景做出不同的技術(shù)決策。7. 訓(xùn)練過程中遇到的經(jīng)典問題排查實(shí)錄7.1 損失不下降我接觸的初學(xué)者項目里損失不下降或下降極慢是最常見的問題。排查順序應(yīng)該按優(yōu)先級這樣走首先是學(xué)習(xí)率。學(xué)習(xí)率太小相當(dāng)于一個人走路步子邁得極小半天走不出原地。把學(xué)習(xí)率從默認(rèn)的 0.001 調(diào)到 0.01 或 0.1 測試一下。如果損失反而變成 NaN說明學(xué)習(xí)率又太大了中間找平衡點(diǎn)。我常用的粗暴方式是 0.1、0.01、0.001 各試 20 個 epoch看哪個收斂快且穩(wěn)定。其次是數(shù)據(jù)和標(biāo)簽是否對齊。檢查 dataloader 返回的圖片和標(biāo)簽是否匹配我遇到過圖片讀取順序和標(biāo)簽順序錯位的問題導(dǎo)致模型一直在學(xué)錯誤對應(yīng)關(guān)系。可以在 DataLoader 里加調(diào)試代碼打印一批圖片對應(yīng)的標(biāo)簽確認(rèn)。還有一個隱蔽的問題預(yù)處理時沒有做歸一化或歸一化錯誤。如果圖片像素值還是 0~255 范圍模型第一層卷積的輸入范圍就不對特別是使用遷移學(xué)習(xí)模型時預(yù)訓(xùn)練權(quán)重期望的是 0~1 歸一化后帶標(biāo)準(zhǔn)差的輸入輸入范圍錯誤會讓特征提取完全崩潰。7.2 顯存溢出顯存溢出Out of Memory在小顯存機(jī)器上訓(xùn)練大模型時幾乎必現(xiàn)。最有效的排查思路是降低 batch size從 32 降到 16 甚至 8。因為顯存占用主要是每個 batch 的中間特征圖和梯度batch size 直接線性影響顯存。另一個思路是減少圖片輸入尺寸224 降到 160 或 128這對顯存的影響是平方級別的。還有一招是使用梯度累積也就是用多個小 batch 模擬大 batch 的效果。比如顯存只夠跑 batch size 8但你需要 batch size 32那就跑 4 個 batch 后累積梯度再更新一次參數(shù)。7.3 過擬合怎么判斷判斷過擬合不能只看訓(xùn)練集精度因為訓(xùn)練集精度沖到 100% 也很正常。核心指標(biāo)是訓(xùn)練和驗證差距當(dāng)訓(xùn)練精度 95% 以上而驗證精度停留在 75% 左右且驗證損失在下降后又開始回升這就是典型過擬合。緩解手段按優(yōu)先級排序數(shù)據(jù)增強(qiáng)加碼、加 Dropout 比例、縮小模型規(guī)模、提前停止。不要上來就換更大的模型那樣只會過擬合得更嚴(yán)重。我在做項目時首先看的是數(shù)據(jù)增強(qiáng)有沒有做夠數(shù)據(jù)集的多樣性是否覆蓋了應(yīng)用場景的分布。7.4 類別不平衡問題實(shí)際項目中類別不平衡很常見。比如你做一個缺陷檢測99% 的圖片是合格品只有 1% 是有缺陷的。模型不需要學(xué)任何東西全預(yù)測合格就有 99% 精度但這種模型完全沒用。解決辦法一是加權(quán)損失函數(shù)PyTorch 的 CrossEntropyLoss 支持 weight 參數(shù)class_weights torch.tensor([1.0, 10.0]) # 少數(shù)類權(quán)重設(shè)為大的值 criterion nn.CrossEntropyLoss(weightclass_weights.to(device))解決辦法二是采樣策略對少數(shù)類過采樣、對多數(shù)類欠采樣。我自己更傾向于先用加權(quán)損失它簡單直接改動量小實(shí)用性高。7.5 亂加預(yù)訓(xùn)練權(quán)重導(dǎo)致的報錯用遷移學(xué)習(xí)模型時遇到 state_dict 加載報錯一般是新舊模型結(jié)構(gòu)不一致。最常發(fā)生的情況是你用了一個 torchvision 模型然后手動修改了 model.fc但加載預(yù)訓(xùn)練權(quán)重時服務(wù)端期望的還是原始完整結(jié)構(gòu)。解決辦法是設(shè)置 weightsNone 構(gòu)造模型后再加載自己的權(quán)重同時確保模型結(jié)構(gòu)完全一致。8. 實(shí)操總結(jié)與擴(kuò)展方向用 PyTorch 識別圖片這個需求看起來簡單但把一個模型真正跑通、跑好中間的信息量其實(shí)非常大。我最后分享幾個經(jīng)驗數(shù)據(jù)永遠(yuǎn)是第一位的。模型結(jié)構(gòu)再先進(jìn)數(shù)據(jù)沒準(zhǔn)備好效果也出不來。你花在數(shù)據(jù)清洗、數(shù)據(jù)增強(qiáng)上的時間永遠(yuǎn)比調(diào)模型結(jié)構(gòu)參數(shù)更值得。遷移學(xué)習(xí)是現(xiàn)在絕大多數(shù)場景的首選方案。自己從頭訓(xùn)練一個卷積神經(jīng)網(wǎng)絡(luò)除非你有幾十萬上百萬級的標(biāo)注數(shù)據(jù)否則很難超過一個微調(diào)好的預(yù)訓(xùn)練模型。用預(yù)訓(xùn)練模型不是偷懶是把前人幾十年的經(jīng)驗?zāi)眠^來用這是聰明人的選擇。訓(xùn)練日志一定要完整記錄。損失值、精度、學(xué)習(xí)率、模型在哪個 epoch 保存的這些信息看起來不起眼但在你回溯問題、復(fù)現(xiàn)實(shí)驗時能省下大把時間。我自己現(xiàn)在每跑一個實(shí)驗都會把超參數(shù)和日志文件一起歸檔。做完一個基礎(chǔ)分類項目后你可以順這個方向繼續(xù)深入一是多標(biāo)簽分類一張圖中同時識別多個物體二是目標(biāo)檢測不只是分類而是框出物體的位置三是語義分割對每個像素做類別預(yù)測。這三個方向都是在圖像識別基礎(chǔ)上的自然延伸掌握了本文的基礎(chǔ)流程后你會發(fā)現(xiàn)它們很多思維方式是一脈相承的。希望這套完整的流程和代碼能夠幫你把第一個 PyTorch 圖像識別項目順利跑通。在實(shí)際操作中遇到問題歡迎把你的報錯信息和實(shí)驗現(xiàn)象記錄下來大家一起交流排查的思路。 PyTorch圖像分類實(shí)戰(zhàn)：從數(shù)據(jù)處理到遷移學(xué)習(xí)的完整流程
1. 項目概述從一張圖片到識別結(jié)果到底發(fā)生了什么用 PyTorch 做圖片識別核心任務(wù)通常是圖像分類也就是我們輸入一張圖片模型輸出它屬于哪個類別。比如一張貓的照片經(jīng)過模型前向傳播得到貓 0.92、狗 0.05、鳥 0.02這樣的概率分布最后取最高概率對應(yīng)的類別作為預(yù)測結(jié)果。很多初學(xué)者拿到這個需求時第一反應(yīng)是調(diào)個模型跑一下但實(shí)際落地時你會發(fā)現(xiàn)整個流程是一條完整鏈路圖像讀取、預(yù)處理、數(shù)據(jù)加載、模型搭建、訓(xùn)練迭代、驗證調(diào)參、固定權(quán)重、部署推理。任何一個環(huán)節(jié)處理不當(dāng)訓(xùn)練出來的模型精度都會差得離譜。這篇文章我會把這條鏈路每一個環(huán)節(jié)的實(shí)操細(xì)節(jié)都拆開講清楚包括一些你翻文檔翻不到的經(jīng)驗和坑。這篇文章適合誰看剛接觸 PyTorch、想快速跑通一個圖像分類項目的同學(xué)或者已經(jīng)能跑通官方 Demo 但換了自己的數(shù)據(jù)集就不知所措的朋友。我會用貓狗分類作為貫穿全文的示例任務(wù)因為它的數(shù)據(jù)獲取簡單、類別直觀、訓(xùn)練收益明顯。文章里涉及的概念和代碼寫法換成任何其他分類任務(wù)都通用。2. 核心鏈路拆解圖片識別的五段流程2.1 圖像分類的本質(zhì)圖像分類是計算機(jī)視覺最基礎(chǔ)的任務(wù)。它的數(shù)學(xué)本質(zhì)是學(xué)習(xí)一個映射函數(shù) f輸入一張 H×W×C 的圖片張量輸出一個長度為 K 的向量K 是類別總數(shù)向量中的值經(jīng)過 Softmax 歸一化后表示每個類別的概率。在這里有一個很多新手沒想明白的點(diǎn)圖片輸入到網(wǎng)絡(luò)里其實(shí)是一堆數(shù)字。比如一張 224×224 的 RGB 圖片輸入張量的形狀是 [3, 224, 224]取值范圍在 0 到 1 之間。卷積神經(jīng)網(wǎng)絡(luò)在做的事情本質(zhì)上是從這些數(shù)字中逐層提取特征——淺層卷積提取邊緣、顏色、紋理深層卷積提取眼睛、耳朵、輪子這種語義級別的部件。最后一層全連接將這些特征映射到類別空間。所以整個識別流程的五個環(huán)節(jié)是圖像預(yù)處理 → 數(shù)據(jù)加載 → 模型前向傳播 → 損失計算與反向傳播 → 推理預(yù)測。環(huán)境準(zhǔn)備完成后后面每個環(huán)節(jié)我會單獨(dú)展開講。2.2 環(huán)境準(zhǔn)備與硬件選型先說你需要的環(huán)境這里給出一個我用下來比較穩(wěn)定的版本組合PyTorch 2.x CUDA 11.8 或 12.1 Python 3.9~3.11 torchvision 配套版本。嚴(yán)格來說我們不建議追求最新版本穩(wěn)定性和兼容性更重要很多 GPU 驅(qū)動對最新的 CUDA 版本支持滯后跑起來反而出各種詭異報錯。有兩點(diǎn)經(jīng)驗說明一下。第一如果你用的是 NVIDIA 顯卡安裝 PyTorch 之前先去官網(wǎng)根據(jù)你的 CUDA 版本選擇對應(yīng)的安裝命令不要直接用 pip install torch 默認(rèn)裝 CPU 版本不然訓(xùn)練速度會慢幾十倍。第二如果沒有獨(dú)立顯卡也可以先跑通整個流程只是訓(xùn)練時間會長很多本文所有代碼在 CPU 上都能運(yùn)行只不過我們建議把 batch size 調(diào)小。3. 數(shù)據(jù)處理與加載比模型更影響精度的地方3.1 數(shù)據(jù)目錄如何組織PyTorch 的 torchvision.datasets.ImageFolder 是處理分類任務(wù)最省事的工具前提是你的數(shù)據(jù)目錄按類別組織成這樣的結(jié)構(gòu)data/ train/ cat/ cat_001.jpg cat_002.jpg dog/ dog_001.jpg dog_002.jpg val/ cat/ cat_001.jpg dog/ dog_001.jpgtrain 和 test 文件夾各自包含類別子文件夾里面的圖片按類別放置。用這種方式組織數(shù)據(jù)ImageFolder 會自動讀取所有圖片并根據(jù)子文件夾名稱生成類別索引我們完全可以用自己真實(shí)的數(shù)據(jù)目錄來替換這個示例結(jié)構(gòu)。一個要注意的點(diǎn)訓(xùn)練集和驗證集必須分開。我見過不少朋友把所有圖片放在一個大文件夾里直接訓(xùn)練這樣驗證集失去了意義。因為模型見過了所有數(shù)據(jù)驗證集精度虛高換到真實(shí)場景性能崩盤。劃分比例一般是 82 或 91。3.2 Dataset 與 DataLoader 的關(guān)鍵配置有了目錄結(jié)構(gòu)接下來就是標(biāo)準(zhǔn)的數(shù)據(jù)加載代碼。這段代碼幾乎是所有 PyTorch 圖像分類項目的骨架你需要根據(jù)自己的數(shù)據(jù)集特點(diǎn)去調(diào)整里面的參數(shù)。import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # 訓(xùn)練集預(yù)處理 train_transform transforms.Compose([ transforms.Resize((256, 256)), # 先縮放 transforms.RandomCrop(224), # 再隨機(jī)裁剪 transforms.RandomHorizontalFlip(), # 隨機(jī)水平翻轉(zhuǎn) transforms.ToTensor(), # 轉(zhuǎn)為張量并歸一化到 [0,1] transforms.Normalize( mean[0.485, 0.456, 0.406], # ImageNet 均值 std[0.229, 0.224, 0.225] # ImageNet 標(biāo)準(zhǔn)差 ) ]) # 驗證集預(yù)處理只做縮放和歸一化不做數(shù)據(jù)增強(qiáng) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) train_dataset datasets.ImageFolder(data/train, transformtrain_transform) val_dataset datasets.ImageFolder(data/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue)這里面的重點(diǎn)在于為什么訓(xùn)練要用縮放 隨機(jī)裁剪而驗證集直接縮放。訓(xùn)練集需要數(shù)據(jù)增強(qiáng)來增加樣本多樣性隨機(jī)裁剪相當(dāng)于在 256×256 的圖片中隨機(jī)截取 224×224 的區(qū)域相當(dāng)于一個樣本生成了很多不同的子樣本能有效緩解過擬合。而驗證集的目的是評估模型的真實(shí)表現(xiàn)所以只用固定縮放保證一致性。另一個容易忽略的配置是 shuffle 參數(shù)。訓(xùn)練集必須 shuffleTrue因為如果不打亂每個 batch 里都是同一類別的圖片梯度下降的方向會來回震蕩收斂速度慢且不穩(wěn)定。驗證集則不需要 shuffle。關(guān)于 Normalize 的 mean 和 std這一組值來自 ImageNet 數(shù)據(jù)集的統(tǒng)計結(jié)果。我們做遷移學(xué)習(xí)時通常沿用這組值因為它和預(yù)訓(xùn)練模型的輸入分布是匹配的。如果從頭訓(xùn)練自己的模型理論上應(yīng)該重新統(tǒng)計自己數(shù)據(jù)集的均值和標(biāo)準(zhǔn)差但實(shí)際操作中大多數(shù)情況下用了 ImageNet 的數(shù)值也不會差太多。3.3 數(shù)據(jù)增強(qiáng)做多少才算合適數(shù)據(jù)增強(qiáng)是把雙刃劍。太少模型容易過擬合太多會給訓(xùn)練制造不必要的難度。以貓狗分類為例我常用的增強(qiáng)組合按優(yōu)先級排列隨機(jī)水平翻轉(zhuǎn)必選不影響語義隨機(jī)裁剪必選模擬尺度變化隨機(jī)顏色抖動可選增強(qiáng)光照魯棒性隨機(jī)旋轉(zhuǎn)最多 ±15 度太大會破壞圖像結(jié)構(gòu)如果數(shù)據(jù)集只有幾千張圖片我建議增強(qiáng)力度適當(dāng)加強(qiáng)把 RandomResizedCrop 的尺度范圍設(shè)為 (0.5, 1.0)讓模型學(xué)會從不同尺度下識別目標(biāo)。如果數(shù)據(jù)集有十萬級別增強(qiáng)可以適度保守重點(diǎn)放在防止過擬合上。我做項目時驗證過一個經(jīng)驗當(dāng)訓(xùn)練集樣本量少于 5000 時數(shù)據(jù)增強(qiáng)的收益是顯著的分類精度可以提升 2 到 5 個百分點(diǎn)。但如果你用的樣本量已經(jīng)很大或者數(shù)據(jù)本身質(zhì)量很高盲目加增強(qiáng)反而可能導(dǎo)致精度不升反降。4. 模型如何選手寫 CNN 還是遷移學(xué)習(xí)4.1 從零搭建一個可用的 CNN如果從頭搭建一個卷積神經(jīng)網(wǎng)絡(luò)核心是感受野的設(shè)計。簡單來說每一層卷積核覆蓋的區(qū)域大小決定了網(wǎng)絡(luò)能看到多大范圍的圖像信息。一個經(jīng)典的小型 CNN 結(jié)構(gòu)如下import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x這個模型的思路是輸入 3×224×224 圖片經(jīng)過三層卷積 池化后特征圖變成 128×28×28然后通過 AdaptiveAvgPool2d 把特征圖壓縮成 1×1再接全連接分類。三層卷積提取不同層次的特征池化逐步降低分辨率擴(kuò)大感受野。從零搭建 CNN 的優(yōu)勢是不依賴外部權(quán)重部署時自由度很大訓(xùn)練過程完全可控。缺點(diǎn)是精度天花板明顯對于復(fù)雜任務(wù)很難超過預(yù)訓(xùn)練模型。所以我的建議是如果數(shù)據(jù)集很小但圖像特征明顯比如純色背景下區(qū)分圓形和方形自建 CNN 足夠用如果圖像復(fù)雜、類別相近優(yōu)先考慮遷移學(xué)習(xí)。4.2 遷移學(xué)習(xí)用預(yù)訓(xùn)練模型大幅提升精度遷移學(xué)習(xí)的邏輯和自建模型完全不同。它的核心是復(fù)用在大規(guī)模數(shù)據(jù)集上訓(xùn)練好的模型參數(shù)——比如在 ImageNet 上已經(jīng)學(xué)好了邊緣、紋理、形狀等通用特征——然后在自己數(shù)據(jù)集上微調(diào)。這就像你不需要重新學(xué)習(xí)認(rèn)識顏色和輪廓只需要學(xué)習(xí)貓和狗的區(qū)別就行。使用 torchvision 的預(yù)訓(xùn)練模型最常用的幾個選項ResNet18 / ResNet50精度和速度最均衡的選擇EfficientNet-B0計算量更小精度相當(dāng)MobileNetV3適合移動端部署場景VGG16結(jié)構(gòu)簡單但參數(shù)量大訓(xùn)練慢實(shí)用性下降以 ResNet18 為例遷移學(xué)習(xí)的標(biāo)準(zhǔn)寫法是凍結(jié)特征層參數(shù)替換最后一層全連接from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 凍結(jié)所有層的參數(shù) for param in model.parameters(): param.requires_grad False # 獲取最后一層全連接的輸入維度 num_features model.fc.in_features # 替換成自己的分類器 model.fc nn.Sequential( nn.Linear(num_features, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, 2) )這段代碼里最關(guān)鍵的邏輯是 requires_grad False。凍結(jié)特征層參數(shù)后反向傳播的梯度不會更新這些層的權(quán)重模型只會訓(xùn)練我們新加的全連接層。這樣有幾個好處一是訓(xùn)練速度大幅提升因為不需要計算大量參數(shù)的梯度二是可以有效防止數(shù)據(jù)集過小導(dǎo)致的過擬合避免破壞預(yù)訓(xùn)練模型已經(jīng)學(xué)好的通用特征表示。微調(diào)有兩種常用策略一是上面這種只訓(xùn)練分類頭的方案適合數(shù)據(jù)集幾百到幾千張的情況二是全量微調(diào)也就是所有層都放開訓(xùn)練適合數(shù)據(jù)集規(guī)模較大時有足夠的樣本來刷新預(yù)訓(xùn)練知識。實(shí)際項目中我建議先用凍結(jié)方式訓(xùn)練幾個 epoch 讓分類頭收斂然后解凍最后幾層做微調(diào)這樣精度更高也穩(wěn)定。4.3 損失函數(shù)和優(yōu)化器的選擇邏輯分類任務(wù)的標(biāo)準(zhǔn)配置是交叉熵?fù)p失函數(shù)加 SGD 或 Adam。這里的交叉熵?fù)p失函數(shù)值得多說一句PyTorch 的 nn.CrossEntropyLoss 其實(shí)做了兩件事——對模型的輸出做 Softmax 歸一化然后計算負(fù)對數(shù)似然損失。所以模型最后一層不需要再接 Softmax直接輸出原始 logits 就好否則等于做了兩次歸一化精度會受影響。criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.fc.parameters(), lr0.001)優(yōu)化器方面Adam 適合快速迭代和調(diào)參對學(xué)習(xí)率的敏感度低一些。SGD 加 Momentum 需要更精細(xì)的學(xué)習(xí)率調(diào)整但收斂后的精度往往更高。我在實(shí)際訓(xùn)練中習(xí)慣先用 Adam 跑通流程之后如果有優(yōu)化需求再切換到帶學(xué)習(xí)率衰減的 SGD。這里有個新手常踩的坑如果只訓(xùn)練分類頭優(yōu)化的參數(shù)只有 model.fc.parameters()因為被凍結(jié)的層沒有梯度。如果全量微調(diào)則要傳入 model.parameters()。如果混用凍結(jié)層和可訓(xùn)練層需要確保優(yōu)化器中傳入了正確的參數(shù)組否則會出現(xiàn)模型一點(diǎn)也不學(xué)習(xí)的困惑現(xiàn)象。4.4 設(shè)備準(zhǔn)備讓代碼自動跑在 GPU 或 CPU 上模型訓(xùn)練前要先確定計算設(shè)備。建議直接寫成下面這種自適應(yīng)方式避免在 GPU 機(jī)器上調(diào)試過的代碼挪到 CPU 機(jī)器上直接報錯device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)model.to(device) 這一步很多人會漏掉導(dǎo)致前向傳播時模型還在 CPU 上數(shù)據(jù)在 GPU 上報錯說 device mismatch。還有一點(diǎn)需要特別注意輸入圖片張量也要調(diào)用 .to(device) 移動設(shè)備我在初學(xué)階段經(jīng)常因為忘記把 data 移動到 GPU 上而報錯這種問題排查起來很浪費(fèi)時間。5. 訓(xùn)練循環(huán)與模型驗證的完整實(shí)現(xiàn)5.1 標(biāo)準(zhǔn)訓(xùn)練循環(huán)代碼訓(xùn)練循環(huán)是整個項目最核心的代碼段。我把一個完整的訓(xùn)練驗證循環(huán)貼出來這份代碼可以直接用以下方式運(yùn)行然后根據(jù)打印的日志判斷訓(xùn)練效果。def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc def validate(model, val_loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc一個容易犯錯的地方是寫法上訓(xùn)練時必須調(diào)用 model.train()驗證時必須調(diào)用 model.eval()。這兩個模式的區(qū)別在于 BatchNorm 和 Dropout 的行為train 模式下 Dropout 會隨機(jī)丟棄神經(jīng)元BatchNorm 會使用 batch 內(nèi)的統(tǒng)計量eval 模式下 Dropout 關(guān)閉BatchNorm 使用訓(xùn)練階段累積的全局統(tǒng)計量。如果驗證時忘了切換到 eval 模式驗證精度會異常波動并且吃進(jìn)隨機(jī)性。另外驗證階段必須包裹 torch.no_grad()它有兩個用處一是停止梯度追蹤節(jié)省顯存二是避免驗證過程意外修改模型參數(shù)。5.2 模型保存策略只保留最優(yōu)權(quán)重訓(xùn)練過程中的模型保存也是有講究的。不經(jīng)思考地只在最后保存一次模型如果你中途過擬合了最后保存的就是過擬合的模型。我建議采用驗證精度最高時保存的策略best_acc 0.0 for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(...) print(fEpoch {epoch1}: train_loss{train_loss:.4f}, train_acc{train_acc:.4f}, fval_loss{val_loss:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(f - 保存最優(yōu)模型驗證精度 {best_acc:.4f})torch.save 建議只保存 model.state_dict()它只保存模型權(quán)重參數(shù)不攜帶完整的模型結(jié)構(gòu)。原因有兩個一是文件體積更小二是加載時更靈活。你在 state_dict 之外還需要單獨(dú)記錄類別映射關(guān)系否則推理時只有數(shù)字類別沒有類別名稱。這個我后面會講。5.3 學(xué)習(xí)率調(diào)整和 early stopping 的實(shí)際價值訓(xùn)練過程中學(xué)習(xí)率設(shè)置是一個常見痛點(diǎn)。學(xué)習(xí)率太大模型震蕩不收斂太小則訓(xùn)練半天損失紋絲不動。我喜歡用 ReduceLROnPlateau 做驗證精度的自適應(yīng)調(diào)整scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience3 ) for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(...) scheduler.step(val_loss)這里傳入的 modemin 表示當(dāng) val_loss 不再下降時把學(xué)習(xí)率乘以 0.5。patience3 意味著連續(xù) 3 個 epoch 驗證損失沒有明顯下降才觸發(fā)衰減。學(xué)習(xí)率衰減的原理是訓(xùn)練初期需要較大的步長快速靠近最優(yōu)點(diǎn)但接近最優(yōu)點(diǎn)后大步長會在最優(yōu)值附近震蕩此時縮小步長可以讓損失繼續(xù)下降。early stopping 指的是當(dāng)驗證損失連續(xù)多個 epoch 不再下降時提前終止訓(xùn)練。這個和保存最優(yōu)模型配合起來用能有效避免過擬合。但我實(shí)際使用中有一個心得early stopping 的 patience 設(shè)置要根據(jù)數(shù)據(jù)集大小靈活調(diào)整小數(shù)據(jù)集設(shè) 5 個 epoch 比較合理大數(shù)據(jù)集可能需要 10 個甚至更多 epoch 的耐心否則容易在驗證損失的局部極小值處提前停止。6. 推理階段加載模型識別單張新圖片6.1 完整的推理代碼訓(xùn)練完成后真正到識別圖片這個動作時代碼要包含三個步驟加載模型權(quán)重、對輸入圖片做和訓(xùn)練時完全一致的預(yù)處理、前向傳播得到預(yù)測結(jié)果。這份推理代碼我是按照生產(chǎn)環(huán)境的標(biāo)準(zhǔn)寫的可以直接用import torch from PIL import Image from torchvision import transforms def predict_image(image_path, model, device, class_names): # 1. 加載并預(yù)處理圖片 image Image.open(image_path).convert(RGB) # 2. 預(yù)處理與驗證集的預(yù)處理完全一致 preprocessing transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) input_tensor preprocessing(image).unsqueeze(0) # 形狀 [1, 3, 224, 224] input_tensor input_tensor.to(device) # 3. 推理 model.eval() with torch.no_grad(): outputs model(input_tensor) probabilities torch.softmax(outputs, dim1) confidence, predicted_idx torch.max(probabilities, dim1) predicted_class class_names[predicted_idx.item()] return predicted_class, confidence.item()調(diào)用示例# 設(shè)定類別映射0 對應(yīng) cat1 對應(yīng) dog class_names [cat, dog] # 加載訓(xùn)練好的最優(yōu)權(quán)重 model models.resnet18(weightsNone) num_features model.fc.in_features model.fc nn.Sequential(...) # 結(jié)構(gòu)必須和訓(xùn)練時完全一致 model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model model.to(device) result, confidence predict_image(test_photo.jpg, model, device, class_names) print(f識別結(jié)果{result}置信度{confidence:.2%})這里的關(guān)鍵坑在于 load_state_dict 前模型結(jié)構(gòu)必須與訓(xùn)練時完全一致。如果你在訓(xùn)練代碼里改過 model.fc 的結(jié)構(gòu)那么推理代碼里也要做一模一樣的修改否則加載權(quán)重時會報 size mismatch。為了避免這種問題我強(qiáng)烈建議把模型構(gòu)建代碼寫在一個函數(shù)里訓(xùn)練和推理都調(diào)用同一個函數(shù)不要二次復(fù)制。6.2 置信度識別結(jié)果的第二落點(diǎn)很多人做圖片識別只看預(yù)測類別但置信度信息其實(shí)更具價值。置信度是 Softmax 輸出的最大概率值表示為模型對該預(yù)測的自信程度。我在實(shí)際項目中總結(jié)的置信度決策規(guī)則置信度 0.9模型非常確定可以放心采納置信度在 0.6 到 0.9 之間模型有一定把握可以采納但要提示用戶核實(shí)置信度低于 0.6建議返回?zé)o法確定因為多類別中每個類別的概率都比較分散強(qiáng)取最大值很可能是錯誤判斷一個常見場景是模型對貓的置信度是 0.45對狗的置信度是 0.42此時取貓作為結(jié)果并不穩(wěn)妥。這種低置信度通常意味著圖片里的目標(biāo)不明顯、遮擋嚴(yán)重或包含訓(xùn)練數(shù)據(jù)中未出現(xiàn)的背景干擾。如果業(yè)務(wù)場景對誤判容忍度低建議加一個置信度閾值判斷。6.3 推理性能優(yōu)化經(jīng)驗如果識別功能要集成到實(shí)際應(yīng)用中推理速度和延遲就很重要了。幾個實(shí)用的優(yōu)化方向把整個模型用 model.eval() 切換去掉 Dropout 的隨機(jī)性對輸入圖片做 batch 推理把多張圖片合成一個 batch 一次前向傳播比單張循環(huán)快很多倍使用 torch.no_grad() 防止構(gòu)建計算圖如果需要更高吞吐量可以考慮用 ONNX 導(dǎo)出模型加速但在精度要求高時要注意導(dǎo)出和推理的精度差異我參與過的一個圖片識別的技術(shù) Demo用 ResNet18 在 GPU 上單張推理延遲只有幾毫秒完全滿足實(shí)時應(yīng)用的需求。而如果換用 CPU 推理延遲會上升到幾十毫秒這種差距會讓一些對時延敏感的應(yīng)用場景做出不同的技術(shù)決策。7. 訓(xùn)練過程中遇到的經(jīng)典問題排查實(shí)錄7.1 損失不下降我接觸的初學(xué)者項目里損失不下降或下降極慢是最常見的問題。排查順序應(yīng)該按優(yōu)先級這樣走首先是學(xué)習(xí)率。學(xué)習(xí)率太小相當(dāng)于一個人走路步子邁得極小半天走不出原地。把學(xué)習(xí)率從默認(rèn)的 0.001 調(diào)到 0.01 或 0.1 測試一下。如果損失反而變成 NaN說明學(xué)習(xí)率又太大了中間找平衡點(diǎn)。我常用的粗暴方式是 0.1、0.01、0.001 各試 20 個 epoch看哪個收斂快且穩(wěn)定。其次是數(shù)據(jù)和標(biāo)簽是否對齊。檢查 dataloader 返回的圖片和標(biāo)簽是否匹配我遇到過圖片讀取順序和標(biāo)簽順序錯位的問題導(dǎo)致模型一直在學(xué)錯誤對應(yīng)關(guān)系。可以在 DataLoader 里加調(diào)試代碼打印一批圖片對應(yīng)的標(biāo)簽確認(rèn)。還有一個隱蔽的問題預(yù)處理時沒有做歸一化或歸一化錯誤。如果圖片像素值還是 0~255 范圍模型第一層卷積的輸入范圍就不對特別是使用遷移學(xué)習(xí)模型時預(yù)訓(xùn)練權(quán)重期望的是 0~1 歸一化后帶標(biāo)準(zhǔn)差的輸入輸入范圍錯誤會讓特征提取完全崩潰。7.2 顯存溢出顯存溢出Out of Memory在小顯存機(jī)器上訓(xùn)練大模型時幾乎必現(xiàn)。最有效的排查思路是降低 batch size從 32 降到 16 甚至 8。因為顯存占用主要是每個 batch 的中間特征圖和梯度batch size 直接線性影響顯存。另一個思路是減少圖片輸入尺寸224 降到 160 或 128這對顯存的影響是平方級別的。還有一招是使用梯度累積也就是用多個小 batch 模擬大 batch 的效果。比如顯存只夠跑 batch size 8但你需要 batch size 32那就跑 4 個 batch 后累積梯度再更新一次參數(shù)。7.3 過擬合怎么判斷判斷過擬合不能只看訓(xùn)練集精度因為訓(xùn)練集精度沖到 100% 也很正常。核心指標(biāo)是訓(xùn)練和驗證差距當(dāng)訓(xùn)練精度 95% 以上而驗證精度停留在 75% 左右且驗證損失在下降后又開始回升這就是典型過擬合。緩解手段按優(yōu)先級排序數(shù)據(jù)增強(qiáng)加碼、加 Dropout 比例、縮小模型規(guī)模、提前停止。不要上來就換更大的模型那樣只會過擬合得更嚴(yán)重。我在做項目時首先看的是數(shù)據(jù)增強(qiáng)有沒有做夠數(shù)據(jù)集的多樣性是否覆蓋了應(yīng)用場景的分布。7.4 類別不平衡問題實(shí)際項目中類別不平衡很常見。比如你做一個缺陷檢測99% 的圖片是合格品只有 1% 是有缺陷的。模型不需要學(xué)任何東西全預(yù)測合格就有 99% 精度但這種模型完全沒用。解決辦法一是加權(quán)損失函數(shù)PyTorch 的 CrossEntropyLoss 支持 weight 參數(shù)class_weights torch.tensor([1.0, 10.0]) # 少數(shù)類權(quán)重設(shè)為大的值 criterion nn.CrossEntropyLoss(weightclass_weights.to(device))解決辦法二是采樣策略對少數(shù)類過采樣、對多數(shù)類欠采樣。我自己更傾向于先用加權(quán)損失它簡單直接改動量小實(shí)用性高。7.5 亂加預(yù)訓(xùn)練權(quán)重導(dǎo)致的報錯用遷移學(xué)習(xí)模型時遇到 state_dict 加載報錯一般是新舊模型結(jié)構(gòu)不一致。最常發(fā)生的情況是你用了一個 torchvision 模型然后手動修改了 model.fc但加載預(yù)訓(xùn)練權(quán)重時服務(wù)端期望的還是原始完整結(jié)構(gòu)。解決辦法是設(shè)置 weightsNone 構(gòu)造模型后再加載自己的權(quán)重同時確保模型結(jié)構(gòu)完全一致。8. 實(shí)操總結(jié)與擴(kuò)展方向用 PyTorch 識別圖片這個需求看起來簡單但把一個模型真正跑通、跑好中間的信息量其實(shí)非常大。我最后分享幾個經(jīng)驗數(shù)據(jù)永遠(yuǎn)是第一位的。模型結(jié)構(gòu)再先進(jìn)數(shù)據(jù)沒準(zhǔn)備好效果也出不來。你花在數(shù)據(jù)清洗、數(shù)據(jù)增強(qiáng)上的時間永遠(yuǎn)比調(diào)模型結(jié)構(gòu)參數(shù)更值得。遷移學(xué)習(xí)是現(xiàn)在絕大多數(shù)場景的首選方案。自己從頭訓(xùn)練一個卷積神經(jīng)網(wǎng)絡(luò)除非你有幾十萬上百萬級的標(biāo)注數(shù)據(jù)否則很難超過一個微調(diào)好的預(yù)訓(xùn)練模型。用預(yù)訓(xùn)練模型不是偷懶是把前人幾十年的經(jīng)驗?zāi)眠^來用這是聰明人的選擇。訓(xùn)練日志一定要完整記錄。損失值、精度、學(xué)習(xí)率、模型在哪個 epoch 保存的這些信息看起來不起眼但在你回溯問題、復(fù)現(xiàn)實(shí)驗時能省下大把時間。我自己現(xiàn)在每跑一個實(shí)驗都會把超參數(shù)和日志文件一起歸檔。做完一個基礎(chǔ)分類項目后你可以順這個方向繼續(xù)深入一是多標(biāo)簽分類一張圖中同時識別多個物體二是目標(biāo)檢測不只是分類而是框出物體的位置三是語義分割對每個像素做類別預(yù)測。這三個方向都是在圖像識別基礎(chǔ)上的自然延伸掌握了本文的基礎(chǔ)流程后你會發(fā)現(xiàn)它們很多思維方式是一脈相承的。希望這套完整的流程和代碼能夠幫你把第一個 PyTorch 圖像識別項目順利跑通。在實(shí)際操作中遇到問題歡迎把你的報錯信息和實(shí)驗現(xiàn)象記錄下來大家一起交流排查的思路。