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