從零搭一個食物圖片分類器我的 PyTorch 入門實(shí)戰(zhàn)記錄想入門 PyTorch但又不想只跑個 MNIST 就完事那食物圖片分類是個特別合適的練手任務(wù)。它既不像 MNIST 那樣簡單到?jīng)]有成就感也不像 COCO 檢測那樣一上來就需要分布式訓(xùn)練和多卡策略難度剛好卡在一個“跳一跳夠得著”的位置。這篇文章是我自己從零開始實(shí)現(xiàn)一個食物圖片分類項目的完整記錄包含數(shù)據(jù)集整理、預(yù)處理細(xì)節(jié)、模型選型、訓(xùn)練循環(huán)手寫、調(diào)參經(jīng)驗、踩坑復(fù)盤。項目本身不復(fù)雜但能把 PyTorch 的日常工作流全部過一遍。做完之后你會發(fā)現(xiàn)后面再去碰遷移學(xué)習(xí)、目標(biāo)檢測、語義分割這些任務(wù)底子已經(jīng)打得差不多了。適合誰看剛學(xué)完 Python 基礎(chǔ)、看過幾篇 PyTorch 教程但沒完整做過項目的朋友以及想用一份清晰代碼把“訓(xùn)練-評估-預(yù)測-模型保存”整條鏈路跑通的人。如果你是資深煉丹師這篇文章可能偏基礎(chǔ)但里面有幾個關(guān)于數(shù)據(jù)加載和調(diào)參的坑也許能讓會心一笑。1. 項目整體拆解食物分類到底在解決什么問題1.1 食物分類任務(wù)的核心難點(diǎn)食物圖片分類本質(zhì)上是一個圖像分類問題模型要做的就是輸入一張圖片輸出它屬于哪個類別。我這邊定了 10 個常見食物類別面包、乳酪、冰淇淋、蘋果、香蕉、玉米、胡蘿卜、煎蛋、米飯、意面。聽起來和 CIFAR-10 很像但實(shí)際做起來有幾個地方和普通分類不一樣。第一是類內(nèi)差異大。同樣是米飯有炒飯、白米飯、蓋澆飯顏色、紋理、擺盤方式完全不同。同一個類別里的圖片可能比不同類別之間的差異還大這對模型的特征提取能力提出了要求。第二是類間相似度高。胡蘿卜和玉米都是長條形的蘋果和番茄都是圓形的僅僅靠顏色或者形狀一個維度的特征根本分不開。這也是為什么傳統(tǒng)手工特征做食物分類效果一直不好的原因需要模型自己去學(xué)更抽象的組合特征。第三是數(shù)據(jù)噪聲多。食物圖片來自真實(shí)生活場景有餐盤、桌面、餐具、手指、光影干擾甚至有咬過一口的蘋果。如果只用干凈的商品圖訓(xùn)練到了真實(shí)場景基本就崩了所以數(shù)據(jù)增強(qiáng)在這個任務(wù)里不是可選項而是必需品。1.2 為什么選擇 PyTorch 而不是其他框架做這個項目我選 PyTorch不是因為別的框架不行而是因為 PyTorch 的調(diào)試體驗對新手最友好。最核心的一點(diǎn)是動態(tài)計算圖。模型跑起來之后你可以隨時在任意一行代碼處打斷點(diǎn)查看中間張量的 shape、dtype、梯度值改完代碼不用重新構(gòu)建整個圖就能繼續(xù)跑。這種交互式的開發(fā)體驗對新手理解“前向傳播-反向傳播”的整個過程特別有幫助。你自己手寫一個訓(xùn)練循環(huán)比用任何高層封裝都更能理解框架內(nèi)部發(fā)生了什么。另外 PyTorch 的生態(tài)在視覺領(lǐng)域是最完整的。torchvision 模型庫、transform 工具、預(yù)訓(xùn)練權(quán)重、各類第三方擴(kuò)展基本上你能想到的視覺處理需求都有現(xiàn)成的輪子。社區(qū)討論也多遇到問題搜一下大概率能找到答案。1.3 整體技術(shù)選型與管線設(shè)計整個項目按下面這條管線走數(shù)據(jù)采集與整理 - 預(yù)處理 Pipeline - DataLoader 構(gòu)建 - 模型選型 - 訓(xùn)練循環(huán)實(shí)現(xiàn) - 評估與保存 - 單張圖片推理硬件方面我用的是一塊顯存 8GB 的中端顯卡。這個配置訓(xùn)練 10 分類的 ResNet18 完全夠用Batch Size 開到 64 也沒壓力。沒有 GPU 也能跑就是慢一些CPU 模式下建議用更小的模型或者把圖片尺寸降下來。軟件環(huán)境方面PyTorch 版本用的 2.1Python 3.10torchvision 0.16。框架內(nèi)部的接口這兩年有些變化我代碼里會加一點(diǎn)兼容性說明如果你用的版本不一樣大概率也能對上。2. 數(shù)據(jù)集準(zhǔn)備與預(yù)處理這步做不好后面全白搭2.1 數(shù)據(jù)集來源與整理方式數(shù)據(jù)是分類任務(wù)的命根子。我沒有直接用某個現(xiàn)成的食物數(shù)據(jù)集而是從一個開源圖片平臺上抓取標(biāo)注好的食物圖片再結(jié)合一部分自己手拍的圖片做補(bǔ)充。整理成下面這樣的目錄結(jié)構(gòu)food_dataset/ ├── train/ │ ├── bread/ │ ├── cheese/ │ ├── ice_cream/ │ ├── apple/ │ ├── banana/ │ ├── corn/ │ ├── carrot/ │ ├── fried_egg/ │ ├── rice/ │ └── spaghetti/ └── val/ ├── bread/ └── ...每個類別我收集了大約 300 張訓(xùn)練圖和 80 張驗證圖。這個數(shù)據(jù)量對于初學(xué)者完全夠用了如果再少可能需要更強(qiáng)的數(shù)據(jù)增強(qiáng)或遷移學(xué)習(xí)來兜底。目錄按類別分好之后用 torchvision.datasets.ImageFolder 就可以直接讀取它會把子目錄名自動映射成類別標(biāo)簽。這里有個關(guān)鍵點(diǎn)子目錄的命名盡量用英文避免編碼問題導(dǎo)致讀取錯誤因為某些系統(tǒng)對中文路徑默認(rèn)編碼是 GBKPyTorch 底層用的是 UTF-8碰到中文目錄名容易臨時找不到文件。2.2 預(yù)處理 Pipeline 的幾個關(guān)鍵細(xì)節(jié)預(yù)處理是整個項目最容易忽略但影響最大的部分。torchvision 的 transform 流水線我這樣定義from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.6, 1.0)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])訓(xùn)練集用了 RandomResizedCrop這一步很關(guān)鍵它不僅僅是裁剪還會隨機(jī)改變圖片的縮放比例和寬高比相當(dāng)于做了模擬距離變化和拍攝角度的增強(qiáng)。scale 參數(shù)我設(shè)成 (0.6, 1.0)意思是裁剪區(qū)域占原圖的 60% 到 100%。這個區(qū)間對食物比較合適太低了把整個食物裁掉一半反而引入無效噪聲。歸一化的 mean 和 std 直接用 ImageNet 的統(tǒng)計值因為我們后面會加載 ImageNet 預(yù)訓(xùn)練權(quán)重輸入分布必須和預(yù)訓(xùn)練時一致否則前期微調(diào)階段梯度會不穩(wěn)定收斂慢很多。驗證集一定不要用 RandomResizedCrop。很多新手在驗證集也上隨機(jī)裁剪導(dǎo)致每次評估同一張圖都可能得到不同的預(yù)測結(jié)果最后模型得分忽高忽低根本沒法判斷好壞。驗證集用 Resize 到 256 再 CenterCrop 到 224這個流程是 ImageNet 分類任務(wù)的標(biāo)配好處是既保留中心主體信息又統(tǒng)一了輸入尺寸。2.3 數(shù)據(jù)劃分與 DataLoader 的優(yōu)化配置數(shù)據(jù)劃分我是直接在目錄結(jié)構(gòu)層面完成的訓(xùn)練和驗證獨(dú)立文件夾簡單可靠。沒有做 K 折交叉驗證因為對初學(xué)者來說單次劃分配合早停法已經(jīng)足夠判斷模型好壞。DataLoader 的配置里有幾個參數(shù)值得注意from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue)num_workers 設(shè)為 4在我的機(jī)器上數(shù)據(jù)加載和訓(xùn)練基本可以并行GPU 不會因為等數(shù)據(jù)而空轉(zhuǎn)。這個參數(shù)不是越大越好太多進(jìn)程反而會增加 CPU 調(diào)度開銷和內(nèi)存占用。經(jīng)驗值一般是 CPU 核心數(shù)的一半左右如果你的機(jī)器是 8 核設(shè) 4 是安全的起點(diǎn)。pin_memoryTrue 也很重要它把數(shù)據(jù)加載進(jìn)鎖頁內(nèi)存從 CPU 拷貝到 GPU 時會走更快的 DMA 通道。代價是占用一部分固定內(nèi)存但性能收益在數(shù)據(jù)量大的時候非常明顯。這里有個小坑如果你的內(nèi)存只有 8GB數(shù)據(jù)又特別大的情況下pin_memory 可能導(dǎo)致內(nèi)存壓力過大那就別開。3. 模型構(gòu)建從手寫網(wǎng)絡(luò)到 ResNet 遷移學(xué)習(xí)3.1 模型結(jié)構(gòu)選型與對比模型選型我認(rèn)真對比過三個方案。第一個方案是自己寫一個三層卷積網(wǎng)絡(luò)結(jié)構(gòu)像這樣import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes10): 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.Flatten(), nn.Linear(128 * 28 * 28, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, x): return self.classifier(self.features(x))這個網(wǎng)絡(luò)在 CIFAR-10 這種小圖上可以跑一跑但換到 224x224 的食物圖片上三層卷積的視野范圍太小提取不到足夠抽象的語義特征訓(xùn)練兩輪就能看到驗證準(zhǔn)確率卡在 60% 左右上不去。這就是典型的模型容量不足。第二個方案是用 torchvision 里現(xiàn)成的 ResNet18from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, 10)這里最關(guān)鍵的一行是 model.fc nn.Linear(num_ftrs, 10)它把 ResNet18 原來的 1000 類分類頭換成了我們自己的 10 類分類頭。前面的卷積層依然保留原來的參數(shù)這些參數(shù)已經(jīng)在 ImageNet 上學(xué)會了通用的邊緣、紋理、形狀特征相當(dāng)于省掉了從頭學(xué)特征的時間。第三個方案是 ResNet34。比 ResNet18 多了一層殘差塊參數(shù)量大一些但在這個 10 分類小數(shù)據(jù)集上提升不明顯反而訓(xùn)練速度更慢。最終我選擇 ResNet18這個選擇非常穩(wěn)妥預(yù)訓(xùn)練權(quán)重齊全、顯存占用小、推理速度快后續(xù)還可以拿它做實(shí)時 demo。3.2 遷移學(xué)習(xí)的三層境界用預(yù)訓(xùn)練模型做遷移學(xué)習(xí)不同階段策略不一樣。最簡單的做法是凍結(jié)卷積層只訓(xùn)練分類頭。具體操作是將前面所有參數(shù)的 requires_grad 設(shè)為 False只有最后一層保持 Truefor param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True這樣做的好處是訓(xùn)練速度極快因為反向傳播只經(jīng)過最后一層梯度計算量小。缺點(diǎn)是特征提取網(wǎng)絡(luò)是固定的無法針對食物數(shù)據(jù)做自適應(yīng)優(yōu)化。如果數(shù)據(jù)量少這個方法很管用我在 300 張每類的數(shù)據(jù)集上試過驗證準(zhǔn)確率大概能到 75%。更常用的做法是解凍部分網(wǎng)絡(luò)做全參數(shù)微調(diào)。做法就是什么都不用改全網(wǎng)絡(luò)參數(shù)都保持可訓(xùn)練但把學(xué)習(xí)率調(diào)小。因為如果從頭到尾都用默認(rèn)學(xué)習(xí)率 0.01 訓(xùn)練預(yù)訓(xùn)練模型前面卷積層的預(yù)訓(xùn)練特征會被破壞。我用的是不同層不同學(xué)習(xí)率的方法fc_params model.fc.parameters() conv_params [param for name, param in model.named_parameters() if name not in {fc.weight, fc.bias}] optimizer torch.optim.SGD([ {params: conv_params, lr: 1e-4}, {params: fc_params, lr: 1e-3} ], momentum0.9)這樣卷積層用一個較小的學(xué)習(xí)率微調(diào)分類頭用一個較大的學(xué)習(xí)率從頭學(xué)。經(jīng)驗證明這樣做比全局統(tǒng)一學(xué)習(xí)率效果更好能兼顧特征適配和分類器收斂速度。還有一個進(jìn)階做法叫層級解凍就是把網(wǎng)絡(luò)從后往前逐步放開訓(xùn)練每訓(xùn)練幾個 epoch 解凍一層。實(shí)際操作中感覺收益性價比不高在數(shù)據(jù)集不太大的情況下全參數(shù)微調(diào) 小學(xué)習(xí)率就已經(jīng)很好了。3.3 損失函數(shù)、優(yōu)化器與評估指標(biāo)的選擇分類任務(wù)最常用的損失函數(shù)是交叉熵?fù)p失criterion nn.CrossEntropyLoss()這個函數(shù)把模型的原始輸出 logits 和標(biāo)簽作為輸入內(nèi)部自動完成 softmax 計算和交叉熵運(yùn)算。需要注意模型輸出的是 logits不要在外面再包一層 softmax。新手經(jīng)常犯這個錯誤導(dǎo)致訓(xùn)練初期損失值異常地高因為 softmax 之后再算交叉熵數(shù)值穩(wěn)定性會變差。如果需要在推理時得到概率輸出那是在模型輸出之后單獨(dú)調(diào) softmax而不是加在模型結(jié)構(gòu)里。優(yōu)化器我最終選了 SGD 而不是 Adam。很多人覺得 Adam 好收斂快但是 SGD 在配合動量之后泛化能力往往比 Adam 好尤其在數(shù)據(jù)量不大的情況下。學(xué)習(xí)率我設(shè)成 0.001 起步配合余弦退火調(diào)度器。評估指標(biāo)不能只看準(zhǔn)確率。食物分類的類別不均衡問題不嚴(yán)重但為了更穩(wěn)妥我還是同時監(jiān)控每個類別的準(zhǔn)確率和 Macro F1。準(zhǔn)確率容易給人一種“模型已經(jīng)很好”的錯覺因為如果某幾個類別特別簡單模型會靠它們拉高整體準(zhǔn)確率但真正難的類別可能完全沒學(xué)會。4. 訓(xùn)練循環(huán)手寫與參數(shù)調(diào)優(yōu)模型的成長過程4.1 一個標(biāo)準(zhǔn)訓(xùn)練循環(huán)的完整實(shí)現(xiàn)我堅持手寫訓(xùn)練循環(huán)不用任何高層封裝。原因很簡單訓(xùn)練循環(huán)是 PyTorch 的心臟不手寫一遍就永遠(yuǎn)隔著一層紗。完整代碼如下import torch device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) best_val_acc 0.0 epochs 30 for epoch in range(epochs): model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() train_loss running_loss / len(train_loader) train_acc 100.0 * correct / total # 驗證階段 model.eval() val_loss 0.0 val_correct 0 val_total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) val_loss loss.item() _, predicted torch.max(outputs.data, 1) val_total labels.size(0) val_correct (predicted labels).sum().item() val_acc 100.0 * val_correct / val_total print(fEpoch {epoch1:02d} | Train Loss: {train_loss:.4f} | fTrain Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth)這一段代碼里有幾個細(xì)節(jié)必須說明。model.train() 和 model.eval() 很重要。這兩個模式切換直接影響 BatchNorm 和 Dropout 的行為。train 模式下 BatchNorm 會使用當(dāng)前 batch 的統(tǒng)計量Dropout 會隨機(jī)丟棄神經(jīng)元eval 模式下 BatchNorm 使用累計的全局統(tǒng)計量Dropout 被關(guān)閉。如果不切換到 eval 模式驗證時結(jié)果會不穩(wěn)定這是個常見坑。with torch.no_grad() 讓驗證階段完全不計算梯度。這能省下大量顯存和計算時間。在驗證階段我們只需要前向傳播不需要通過 loss.backward() 更新梯度所以直接開啟無梯度模式。optimizer.zero_grad() 放在前向之前每次迭代都清空梯度。如果不調(diào)用這個梯度會累加在某一步更新的時候用的就是幾個 batch 的累積梯度訓(xùn)練過程會變得極度不穩(wěn)定。PyTorch 不會自動清梯度這一點(diǎn)和有些高層框架不一樣必須手動清。驗證準(zhǔn)確率超過歷史最好就保存模型。我每次 epoch 結(jié)束都檢查一次保存的時候只保留 val_acc 最高的那一次。這種“早停 快照”策略比無條件跑完所有 epoch 更穩(wěn)妥后面加載到的最佳模型不需要再人工挑選。4.2 超參數(shù)調(diào)優(yōu)的完整記錄學(xué)習(xí)率是最關(guān)鍵的超參數(shù)。我做了三次對比實(shí)驗學(xué)習(xí)率優(yōu)化器最終驗證準(zhǔn)確率訓(xùn)練穩(wěn)定性0.1SGD62%波動大后期發(fā)散0.01SGD78%前期震蕩后期穩(wěn)定0.001SGD85%平穩(wěn)上升穩(wěn)定收斂0.001Adam84%收斂最快但后期略波動0.1 這個學(xué)習(xí)率對于遷移學(xué)習(xí)來說太激進(jìn)了。因為預(yù)訓(xùn)練權(quán)重的特征已經(jīng)很好學(xué)習(xí)率太大一步跨過頭就把原有知識破壞了。0.001 是遷移學(xué)習(xí)微調(diào)階段比較安全的選擇。Batch Size 的影響我也對比了。Batch Size顯存占用訓(xùn)練速度穩(wěn)定性16約 1.5GB最慢震蕩劇烈64約 3GB快平穩(wěn)128約 5GB最快有些微震蕩太小 batch 會產(chǎn)生比較大的梯度噪聲訓(xùn)練震蕩明顯。最后我選 64因為它在我的 8GB 顯存卡上留下了充足的余量還可以順便開個并行的驗證程序。優(yōu)化器動量我設(shè)為 0.9這是一個經(jīng)驗上幾乎不需要調(diào)的默認(rèn)值SGD 加動量能在保證梯度方向穩(wěn)定的同時加快收斂速度。4.3 過擬合的判斷與應(yīng)對方法訓(xùn)練過程中我在第 5 個 epoch 開始觀察到訓(xùn)練準(zhǔn)確率接近 90%但驗證準(zhǔn)確率只有 72%。這就是過擬合的典型信號。食物圖片分類數(shù)據(jù)集量不算大模型容量又大過擬合幾乎是必然的。我從兩方面來應(yīng)對。第一是數(shù)據(jù)增強(qiáng)。項目已經(jīng)掛了 RandomResizedCrop、RandomHorizontalFlip、RandomRotation、ColorJitter這些都屬于基礎(chǔ)操作。為了再強(qiáng)化我給訓(xùn)練集加了一個 MixUp 增強(qiáng)import numpy as np def mixup_data(x, y, alpha0.2): if alpha 0: lam np.random.beta(alpha, alpha) else: lam 1 batch_size x.size()[0] index torch.randperm(batch_size).to(device) mixed_x lam * x (1 - lam) * x[index] y_a, y_b y, y[index] return mixed_x, y_a, y_b, lamMixUp 的思路是把兩張圖片按比例混合標(biāo)簽也跟著混合。比如一個蘋果和一張香蕉按 0.7:0.3 混合模型要學(xué)到的是這個混合圖片有 70% 的蘋果和 30% 的香蕉。這樣模型被迫學(xué)習(xí)軟標(biāo)簽決策邊界更加平滑不容易陷入尖銳的過擬合狀態(tài)。第二是Dropout 和權(quán)值衰減。ResNet18 本身已經(jīng)帶了全局平均池化對于防止過擬合有一定幫助。我在分類頭前面加了 Dropout(0.2)并且設(shè)置了 weight_decay1e-4也就是 L2 正則化限制權(quán)重的絕對值大小。加了這兩層保護(hù)后驗證準(zhǔn)確率有 3-4 個百分點(diǎn)的提升。5. 訓(xùn)練過程中踩過的坑與排查實(shí)錄5.1 顯存不足的排查過程第一次跑 ResNet18 加 batch 128大概在訓(xùn)練到第 3 個 epoch 時爆了顯存。報錯信息是 CUDA out of memory。我的排查步驟是這樣的先看是不是常量緩存占用的顯存用 torch.cuda.empty_cache() 清一輪再試確認(rèn)不是緩存問題后把 batch size 降到 64顯存占用從 5.2GB 降到 3GB。問題解決。這里有一個額外的技巧如果只是顯存峰值爆掉可以開啟 PyTorch 2.0 的自動混合精度訓(xùn)練float16 半精度計算會把顯存占用差不多減半。代碼如下from torch.amp import autocast, GradScaler scaler GradScaler() for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with autocast(device_typecuda): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度對我來說不僅是顯存優(yōu)化更是訓(xùn)練加速手段。我那臺顯卡的 FP16 計算速度接近 FP32 的兩倍整個訓(xùn)練時間縮短了接近一半。初學(xué)者不用害怕這個 API 復(fù)雜直接用就行內(nèi)部機(jī)制暫時不需要完全搞懂。5.2 訓(xùn)練不收斂的問題排查有段時間模型 loss 一直停留在 2.3 左右準(zhǔn)確率始終在 35% 上下像是卡住不動。我排查過后發(fā)現(xiàn)問題出在學(xué)習(xí)率預(yù)熱缺失。遷移學(xué)習(xí)微調(diào)時如果一開始就用中等偏大的學(xué)習(xí)率預(yù)訓(xùn)練特征會在早期被大幅擾動導(dǎo)致 loss 飆升后無法恢復(fù)。解決辦法是在前幾個 epoch 用很小的學(xué)習(xí)率做預(yù)熱from torch.optim import lr_scheduler def warmup_lr(epoch, warmup_epochs3, base_lr0.001): if epoch warmup_epochs: return base_lr * (epoch 1) / warmup_epochs return base_lr scheduler lr_scheduler.LambdaLR(optimizer, lr_lambdawarmup_lr)前 3 個 epoch 學(xué)習(xí)率從 0 逐漸升到 0.001之后再正常訓(xùn)練。加了這一層預(yù)熱后模型在第一個 epoch 的 loss 就降到了 1.5 附近訓(xùn)練明顯走上了正軌。另一個不收斂的常見原因是標(biāo)簽問題。我檢查了一遍標(biāo)簽映射發(fā)現(xiàn) ImageFolder 是按文件名排序的目錄名順序和類別數(shù)組順序可能不一致。排查方法是打印前幾個樣本的標(biāo)簽和類別名對齊關(guān)系確認(rèn)無誤后再訓(xùn)練。5.3 數(shù)據(jù)加載慢與 CPU 瓶頸訓(xùn)練過程中 GPU 利用率經(jīng)常掉到 60% 左右看代碼邏輯應(yīng)該是數(shù)據(jù)加載跟不上。我用 num_workers4 解決了這個問題。注意幾個細(xì)節(jié)workers 不是越多越好。我試過 num_workers8GPU 利用率沒明顯提升反而程序啟動變慢偶爾還會出內(nèi)存錯誤。因為我那臺機(jī)器還要跑其他服務(wù)8 個進(jìn)程把 CPU 占滿了。還可以開啟驗證集數(shù)據(jù)的 preload把整個驗證集一次性加載進(jìn)內(nèi)存# 緩存驗證集到顯存減少重復(fù) IO val_images torch.cat([batch[0] for batch in val_loader], dim0).to(device) val_labels torch.cat([batch[1] for batch in val_loader], dim0).to(device)驗證集圖片數(shù)量少幾百張圖占不了多少顯存但節(jié)省每次驗證時反復(fù)從磁盤讀文件的 IO 時間。實(shí)際體驗是每個 epoch 的驗證時間從 15 秒降到了 3 秒。6. 推理部署與結(jié)果分析模型到底學(xué)到了什么6.1 單張圖片推理的完整流程訓(xùn)練完以后我把模型保存成文件然后做了一個簡單的推理腳本。推理過程是獨(dú)立于訓(xùn)練的模擬真實(shí)使用場景。import torch from torchvision import transforms, models from PIL import Image device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsNone) num_ftrs model.fc.in_features model.fc torch.nn.Linear(num_ftrs, 10) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.to(device) model.eval() class_names [bread, cheese, ice_cream, apple, banana, corn, carrot, fried_egg, rice, spaghetti] transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(test_food.jpg).convert(RGB) img_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(img_tensor) probabilities torch.softmax(outputs, dim1) top_prob, top_class torch.topk(probabilities, k3) for i in range(top_prob.size(1)): idx top_class[0][i].item() print(f{class_names[idx]}: {top_prob[0][i].item():.4f})這里有個小坑就是Image.open之后最好強(qiáng)制 convert(RGB)。有些食物圖片是 RGBA 四通道的 PNG如果不轉(zhuǎn)的話輸入張量就是 4 個通道和模型期望的 3 個通道不匹配直接報錯。模型 state_dict 里只保存了參數(shù)沒有保存結(jié)構(gòu)。加載模型時要先創(chuàng)建一個相同的模型結(jié)構(gòu)再 load_state_dict所以我在推理腳本里重新定義了模型。還有一種更省事的方式是保存整個模型對象但不推薦因為模型結(jié)構(gòu)在 PyTorch 升級后可能不兼容。6.2 混淆矩陣分析跑完測試集之后我輸出了一張混淆矩陣圖我就不貼了重點(diǎn)看幾個數(shù)據(jù)點(diǎn)。意面帶最容易和玉米混淆。細(xì)看原因是兩者的顏色接近都是黃棕色系而且食物照片中都有長條形結(jié)構(gòu)。模型沒能學(xué)到“意面是彎曲細(xì)長的面食玉米是一粒一粒的”這種語義級區(qū)分。香蕉和胡蘿卜也非常容易分不清。測試集里胡蘿卜圖片有不少是切成條狀的和香蕉外形很接近。這提醒了我真實(shí)分類場景中“形狀”和“顏色”并不總是可靠特征模型如果能捕捉到“香蕉有弧度”“胡蘿卜根部有漸細(xì)的尖端”效果會更好。通過混淆矩陣我找到了 3 個樣本量不足的難分對有針對性地多收集了這些類別的圖片加入訓(xùn)練集。下次迭代的驗證準(zhǔn)確率從 85% 提升到了 88%。6.3 模型大小與推理性能訓(xùn)練完成之后我對模型的工程指標(biāo)做了記錄。指標(biāo)數(shù)值模型參數(shù)量約 1120 萬模型文件大小約 43MB單張圖片 CPU 推理耗時約 45ms單張圖片 GPU 推理耗時約 3ms顯存占用約 1.2GB對于移動端或嵌入式部署1120 萬參數(shù)量還是偏大。項目如果后續(xù)擴(kuò)展我會考慮蒸餾成一個小型網(wǎng)絡(luò)或者直接用 MobileNetV3。但在“驗證思路、跑通流程”這個階段ResNet18 的體量是完全合理的。7. 經(jīng)驗的總結(jié)做圖像分類項目最容易忽略的幾件事整個項目跑下來我自己有幾點(diǎn)特別深的體會。數(shù)據(jù)整理階段多花一小時訓(xùn)練階段省一周。目錄結(jié)構(gòu)、類別命名、數(shù)據(jù)校驗、樣本分布檢查這些不起眼的準(zhǔn)備決定了后面所有環(huán)節(jié)是否順暢。我一開始沒有認(rèn)真檢查每個類別的圖片數(shù)量訓(xùn)練到中途才發(fā)現(xiàn)有一個類別圖片數(shù)量比其他類別少了一半導(dǎo)致模型嚴(yán)重偏向樣本多的類別又要重新整理數(shù)據(jù)從頭訓(xùn)練。不要迷信模型越復(fù)雜越好。這個 10 分類任務(wù)用 ResNet18 已經(jīng)溢出真正決定模型上限的是數(shù)據(jù)質(zhì)量、增強(qiáng)策略和調(diào)參細(xì)節(jié)。換更大的模型不會帶來想象中的大提升。手寫訓(xùn)練循環(huán)是值得的。雖然 PyTorch Lightning 之類的框架能省掉大量樣板代碼但對我來說手寫循環(huán)帶來的收益是理解層面的——理解了每個張量的變化、每行代碼的作用之后后面用任何高階框架都只是一個“選擇”的問題而不是被框架牽著走。保存模型時多保存一份帶優(yōu)化器狀態(tài)的完整檢查點(diǎn)。我后來想接著上次的訓(xùn)練繼續(xù)跑發(fā)現(xiàn)只保存了模型參數(shù)優(yōu)化器的動量狀態(tài)沒保存結(jié)果后面幾個 epoch 的收斂速度明顯變慢相當(dāng)于丟掉了一部分訓(xùn)練信息。正確做法是每一輪同時保存 optimizer.state_dict 和調(diào)度器狀態(tài)torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_val_acc: best_val_acc, }, checkpoint.pth)恢復(fù)訓(xùn)練時就三行代碼搞定checkpoint torch.load(checkpoint.pth) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict])最后再說一個小技巧是我后來才養(yǎng)成的習(xí)慣每訓(xùn)練一個模型就在最后帶著 test_loader 跑一遍混淆矩陣同時把預(yù)測錯誤的樣本圖片存下來按類別分文件夾放好。這樣做有兩個直接好處一是能一目了然知道困難樣本長什么樣二是給后續(xù)的數(shù)據(jù)增強(qiáng)方向提供了明確依據(jù)。前幾次做項目我就是只盯著準(zhǔn)確率遇到瓶頸根本沒有方向直到養(yǎng)成這個習(xí)慣之后模型的迭代才真正開始“有路線圖”。整個食物分類項目從頭到尾一共花了不到兩周的下班時間卻把 PyTorch 的完整工作流徹底打通了。如果你也是剛?cè)腴T有一定 Python 基礎(chǔ)想找一個有真實(shí)場景又不至于復(fù)雜的項目食物圖片分類是個不錯的選擇。照著這個思路做一遍動手敲一遍代碼調(diào)試一遍報錯你會發(fā)現(xiàn)自己對深度學(xué)習(xí)模型訓(xùn)練的理解一下子變得具體了很多。