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