簡介基于AlexNet的花卉分類識別系統(tǒng)是一套完整的預(yù)訓(xùn)練模型應(yīng)用項(xiàng)目面向Python與深度學(xué)習(xí)入門者、進(jìn)階學(xué)習(xí)者以及正在籌備畢業(yè)設(shè)計(jì)或課程設(shè)計(jì)的同學(xué)可快速用于圖像分類實(shí)踐或項(xiàng)目立項(xiàng)。系統(tǒng)能區(qū)分10種不同類別的花卉訓(xùn)練數(shù)據(jù)來自Kaggle當(dāng)前數(shù)據(jù)集下準(zhǔn)確率約為96%代碼結(jié)構(gòu)清晰便于二次拓展。壓縮包共49個文件以5個Python腳本、40張JPEG樣本與結(jié)果圖像、1個JSON類別配置及1份Markdown說明為主整體大小僅1.67MB結(jié)構(gòu)緊湊。其中包含訓(xùn)練腳本、預(yù)測腳本、模型定義、工具函數(shù)、數(shù)據(jù)集與類別映射、結(jié)果可視化圖片和README說明文檔能夠幫助讀者理清訓(xùn)練、評估與推斷的完整流程并根據(jù)自帶樣例快速驗(yàn)證效果。已有346人學(xué)習(xí)適合用作入門深度學(xué)習(xí)圖像分類的基礎(chǔ)參考也可作為課程設(shè)計(jì)或工程實(shí)訓(xùn)的起點(diǎn)。1. 花卉分類識別系統(tǒng)為什么值得用AlexNet復(fù)現(xiàn)成本與效果邊界花卉分類識別是圖像分類任務(wù)里最適合上手的場景之一而「基于Python的AlexNet花卉分類識別系統(tǒng)」這個組合幾乎是這類小體量任務(wù)的標(biāo)準(zhǔn)答案。AlexNet雖然發(fā)布年頭不短但它的8層結(jié)構(gòu)、6000萬參數(shù)放在幾十類花卉數(shù)據(jù)上依舊夠用訓(xùn)練門檻低、單張推理在CPU上也能到毫秒級還能直接復(fù)用ImageNet預(yù)訓(xùn)練權(quán)重做遷移學(xué)習(xí)。做這件事的人我通常見到兩類一類是想把課程作業(yè)做成完整Demo的學(xué)生另一類是想快速驗(yàn)證花卉識別效果、不打算從零造輪子的開發(fā)者。下面按數(shù)據(jù)準(zhǔn)備、網(wǎng)絡(luò)實(shí)現(xiàn)、訓(xùn)練調(diào)參、踩坑記錄、推理封裝五個環(huán)節(jié)展開全程給可復(fù)現(xiàn)的Python代碼。2. 數(shù)據(jù)集與預(yù)處理把花卉圖片整理成227×227輸入的Python腳本2.1 數(shù)據(jù)集選型公開花卉集與自建數(shù)據(jù)集的取舍公共數(shù)據(jù)集方面Oxford 17 Flowers17個類別每類約80張和Oxford 102 Flowers102類是最經(jīng)典的兩個選擇很多分類論文都以它們?yōu)榛鶞?zhǔn)。第一次做建議先用17類版本類別少、樣本量適中跑完一輪30個epoch在普通消費(fèi)級顯卡上十幾分鐘就夠適合先把整條鏈路跑通。102類的版本留到后面驗(yàn)證模型泛化能力時再用。如果決定自建數(shù)據(jù)集有個問題要想在前面類別怎么定義。同類花在不同光照、不同角度下差異很大而不同品種的菊花在照片里可能長得幾乎一樣。我一般要求每類不少于100張刻意收進(jìn)同株不同角度、晴天和陰天、近景和遠(yuǎn)景這些情況讓模型學(xué)的是花的形態(tài)而不是拍攝環(huán)境。圖源選擇上用自己拍的照片或明確開放許可的資源更穩(wěn)妥避免后續(xù)演示或部署時惹版權(quán)麻煩。2.2 目錄結(jié)構(gòu)與批量整理腳本8:2切分訓(xùn)練集和驗(yàn)證集數(shù)據(jù)目錄統(tǒng)一成「類別名作為文件夾名」的結(jié)構(gòu)這是PyTorch的ImageFolder默認(rèn)支持的組織方式。整理腳本我一般寫成這樣from pathlib import Path import shutil import random src_root Path(raw_flower_photos) # 原始圖片第一層目錄是類別名 train_root Path(flower_data/train) val_root Path(flower_data/val) for label_dir in src_root.iterdir(): if not label_dir.is_dir(): continue # 一次收集該類別下所有常見格式的圖片 images [p for ext in (*.jpg, *.jpeg, *.png) for p in label_dir.rglob(ext)] if len(images) 20: print(f警告{label_dir.name} 只有{len(images)}張圖建議補(bǔ)充) continue random.shuffle(images) split int(len(images) * 0.8) for img in images[:split]: dst train_root / label_dir.name / img.name dst.parent.mkdir(parentsTrue, exist_okTrue) shutil.copy2(img, dst) for img in images[split:]: dst val_root / label_dir.name / img.name dst.parent.mkdir(parentsTrue, exist_okTrue) shutil.copy2(img, dst)邏輯說明遍歷原始目錄下每個類別文件夾把該類別所有jpg、jpeg、png圖片的路徑收集起來隨機(jī)打亂后按8:2比例分別復(fù)制到train和val下對應(yīng)的類別文件夾。用copy2而不是move是為了保留一份原始素材之后補(bǔ)充了新圖想重新切分時不需要重新采集。參數(shù)說明split 0.8是經(jīng)驗(yàn)值。每類樣本不足50張時我會改成0.85盡量多留訓(xùn)練樣本樣本充足時0.7也可以驗(yàn)證集更有統(tǒng)計(jì)意義。腳本里漏掉格式后綴是常見的翻車點(diǎn)如果原始圖里有bmp或webp記得加進(jìn)擴(kuò)展名元組里否則訓(xùn)練集比預(yù)想少一大截。2.3 預(yù)處理管道Resize、Crop與ImageNet歸一化的參數(shù)細(xì)節(jié)AlexNet默認(rèn)輸入是227×227×3的RGB圖注意不是224×224這個數(shù)字經(jīng)常被人抄錯。原因在原始論文第一層卷積核是11×11、步長4輸入227時輸出才是55×55floor((227-11)/4)155。拿224硬套第一層輸出變成54×54后面每層尺寸都會錯位報錯時一堆人盯著維度問題排查半天。既然要用ImageNet預(yù)訓(xùn)練權(quán)重預(yù)處理也得和ImageNet保持一致from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(227), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.4, contrast0.4, saturation0.4), 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, 256)), transforms.CenterCrop(227), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])邏輯說明訓(xùn)練時先Resize到256再RandomResizedCrop到227等于每輪從原圖隨機(jī)位置、隨機(jī)大小裁一塊區(qū)域是一種不需要額外數(shù)據(jù)就能生效的數(shù)據(jù)增強(qiáng)驗(yàn)證時不希望結(jié)果帶隨機(jī)性所以固定用CenterCrop裁正中間。Normalize的均值和標(biāo)準(zhǔn)差取自ImageNet全量統(tǒng)計(jì)只要用預(yù)訓(xùn)練權(quán)重這組值就不能動改了等于把輸入分布整體偏移訓(xùn)練會明顯變慢。參數(shù)說明RandomResizedCrop默認(rèn)scale范圍是(0.08, 1.0)。對花朵這類主體通常占畫面較大面積的場景我會把scale收緊到(0.5, 1.0)避免裁出來的小塊里幾乎沒有花。ColorJitter里saturation建議克制一點(diǎn)花卉識別對顏色敏感飽和度抖動過猛會讓模型學(xué)到失真的顏色分布。3. AlexNet模型搭建用PyTorch逐層實(shí)現(xiàn)卷積與全連接3.1 AlexNet結(jié)構(gòu)拆解從卷積到全連接的層與尺寸AlexNet的主體可以拆成特征提取器加分類器兩段。特征提取器是5層卷積前兩層后面接MaxPool第3、4層卷積不接池化第5層卷積后面再接一次MaxPool分類器是3層全連接前兩層各接一個ReLU和一個Dropout。整個網(wǎng)絡(luò)的核心思想很簡單用卷積逐步把空間尺寸壓小、把通道數(shù)加多最后展平成全連接向量做分類。網(wǎng)上不少實(shí)現(xiàn)把第一層卷積核數(shù)量寫成96這是照抄原始論文的結(jié)果。原始論文因?yàn)楫?dāng)年單卡顯存不夠把96個卷積核拆到兩張GPU上并行訓(xùn)練torchvision的現(xiàn)代單卡實(shí)現(xiàn)第一層只有64個卷積核。這個差異直接決定你能不能加載官方預(yù)訓(xùn)練權(quán)重后面避坑章節(jié)會詳細(xì)展開。各層參數(shù)可以對照下面這張表層名類型卷積核/池化窗口輸出尺寸Conv1卷積64個11×11, stride 455×55×64MaxPool1池化3×3, stride 227×27×64Conv2卷積192個5×5, padding 227×27×192MaxPool2池化3×3, stride 213×13×192Conv3卷積384個3×3, padding 113×13×384Conv4卷積256個3×3, padding 113×13×256Conv5卷積256個3×3, padding 113×13×256MaxPool3池化3×3, stride 26×6×256FC1全連接9216→40964096FC2全連接4096→40964096FC3全連接4096→類別數(shù)類別數(shù)3.2 PyTorch逐層實(shí)現(xiàn)可直接復(fù)制的AlexNet代碼實(shí)際項(xiàng)目里我不會自己從零寫features部分直接用torchvision.models.alexnet接預(yù)訓(xùn)練權(quán)重最省事但為了把結(jié)構(gòu)講清楚先給你一份自建版本import torch.nn as nn class AlexNet(nn.Module): def __init__(self, num_classes10, dropout0.5): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 64, kernel_size11, stride4), # 227 - 55 nn.ReLU(inplaceTrue), nn.MaxPool2d(3, stride2), # 55 - 27 nn.Conv2d(64, 192, kernel_size5, padding2), # 27 - 27 nn.ReLU(inplaceTrue), nn.MaxPool2d(3, stride2), # 27 - 13 nn.Conv2d(192, 384, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(384, 256, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(256, 256, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(3, stride2), # 13 - 6 ) self.classifier nn.Sequential( nn.Dropout(dropout), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplaceTrue), nn.Dropout(dropout), nn.Linear(4096, 4096), nn.ReLU(inplaceTrue), nn.Linear(4096, num_classes), ) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) # 展平保留batch維度 x self.classifier(x) return x邏輯說明features部分把227×227×3的輸入逐步變成6×6×256的特征圖view(x.size(0), -1)把每個樣本展平成9216維向量再進(jìn)分類器。Dropout放在全連接層前面而不是卷積層前面這是AlexNet的原始設(shè)計(jì)——卷積層本身有空間冗余對Dropout不敏感全連接層參數(shù)占比最大、最容易過擬合所以只在分類器里加。參數(shù)說明dropout0.5是原始論文的值在遷移學(xué)習(xí)場景下如果數(shù)據(jù)量不大我會把前面那個Dropout提到0.7后面那個保持0.5。原始論文里還有一個Local Response Normalization層現(xiàn)代框架里它對最終精度的貢獻(xiàn)很小還會拖慢訓(xùn)練和推理速度我一般直接省略實(shí)測對花卉分類這類小類別任務(wù)幾乎沒有影響。3.3 訓(xùn)練前必須定的參數(shù)Dropout、學(xué)習(xí)率與預(yù)訓(xùn)練權(quán)重模型結(jié)構(gòu)定了之后有三個參數(shù)要在寫訓(xùn)練腳本前想清楚。第一個是學(xué)習(xí)率。用預(yù)訓(xùn)練權(quán)重做微調(diào)時我習(xí)慣從0.001起步比從頭訓(xùn)練常用的0.01低一個數(shù)量級因?yàn)轭A(yù)訓(xùn)練權(quán)重已經(jīng)在ImageNet特征空間里收斂得很好了步子邁太大會把學(xué)好的特征沖掉。第二個是優(yōu)化器。花卉分類這種圖像任務(wù)動量SGD仍然比Adam穩(wěn)Adam收斂快但容易落在尖銳的極小值上泛化會差一些。第三個是預(yù)訓(xùn)練權(quán)重的取舍。如果你手里數(shù)據(jù)每類只有幾十張固定住features部分只訓(xùn)練分類器是更穩(wěn)的做法數(shù)據(jù)超過每類200張?jiān)俜砰_全模型微調(diào)。4. 訓(xùn)練與評估讓花卉分類準(zhǔn)確率逼近可用的調(diào)參流程4.1 訓(xùn)練主循環(huán)損失函數(shù)、優(yōu)化器與學(xué)習(xí)率調(diào)度加載數(shù)據(jù)和訓(xùn)練主循環(huán)可以合并成一個腳本關(guān)鍵是把驗(yàn)證邏輯獨(dú)立成函數(shù)每個epoch結(jié)束后跑一次方便觀察曲線。下面這個流程我在多個花卉數(shù)據(jù)集上用的都是同一套骨架import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, models train_ds datasets.ImageFolder(flower_data/train, transformtrain_transform) val_ds datasets.ImageFolder(flower_data/val, transformval_transform) train_loader DataLoader(train_ds, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_ds, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue) device torch.device(cuda if torch.cuda.is_available() else cpu) model models.alexnet(weightsmodels.AlexNet_Weights.IMAGENET1K_V1) in_features model.classifier[6].in_features model.classifier[6] nn.Linear(in_features, len(train_ds.classes)) model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.001, momentum0.9, weight_decay5e-4) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.5) epochs 30 best_acc 0.0 for epoch in range(epochs): model.train() running_loss 0.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() val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch1:02d} | train_loss {running_loss/len(train_loader):.4f} f| val_acc {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_alexnet_flower.pth)邏輯說明models.alexnet(weights...)會自動下載ImageNet預(yù)訓(xùn)練權(quán)重最后一行分類器被替換成你的花卉類別數(shù)量這是遷移學(xué)習(xí)的標(biāo)準(zhǔn)改法。訓(xùn)練循環(huán)里每個batch做一次清零、前向、算損失、反向、更新五步。驗(yàn)證和保存都放在epoch層面做以驗(yàn)證集準(zhǔn)確率為基準(zhǔn)保存最佳權(quán)重而不是用最后一個epoch的權(quán)重。參數(shù)說明batch_size64在8GB顯存上跑227×227輸入沒有壓力顯存不夠就降到32。lr0.001配合StepLR每5個epoch降一半是微調(diào)場景的穩(wěn)妥組合。weight_decay5e-4是AlexNet原始配置全模型微調(diào)時保留它能明顯抑制過擬合。num_workers4在Windows上如果報DataLoader相關(guān)錯誤改成0跑單進(jìn)程最省心。4.2 訓(xùn)練曲線診斷過擬合、欠擬合與Early Stopping訓(xùn)練曲線比任何指標(biāo)都先暴露問題。正常情況應(yīng)該是訓(xùn)練Loss一路下降驗(yàn)證Loss先降后平緩驗(yàn)證準(zhǔn)確率同步爬升。如果訓(xùn)練Loss還在降、驗(yàn)證Loss卻掉頭向上說明模型開始背訓(xùn)練集了這時候繼續(xù)訓(xùn)練只會讓驗(yàn)證集越來越差。Early Stopping的實(shí)現(xiàn)不復(fù)雜也不需要額外裝庫記錄當(dāng)前最佳驗(yàn)證準(zhǔn)確率連續(xù)N個epoch沒有刷新就提前終止。N我一般取7到10。如果你用的是上面那個腳本torch.save已經(jīng)天然只保留最佳權(quán)重所以Early Stopping只是幫你省訓(xùn)練時間不省也可以但數(shù)據(jù)量小的時候多跑那幾個epoch純屬浪費(fèi)電。4.3 評估不止看準(zhǔn)確率混淆矩陣與單類召回率準(zhǔn)確率是平均數(shù)最容易騙人。如果數(shù)據(jù)集里玫瑰占了一半模型把所有圖片都判成玫瑰準(zhǔn)確率也有50%但你真正關(guān)心的可能是郁金香一張都認(rèn)不出來。花卉分類的評估我至少要打印一份分類報告和混淆矩陣from sklearn.metrics import classification_report, confusion_matrix import numpy as np def evaluate(model, loader, criterion, device): model.eval() total, correct, loss_sum 0, 0, 0.0 all_preds, all_labels [], [] with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) loss_sum criterion(outputs, labels).item() _, preds torch.max(outputs, 1) total labels.size(0) correct (preds labels).sum().item() all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) report classification_report( all_labels, all_preds, target_namesloader.dataset.classes) print(report) cm confusion_matrix(all_labels, all_preds) print(混淆矩陣\n, cm) return loss_sum / len(loader), correct / total邏輯說明torch.no_grad()里面只做前向不做反向省顯存也提速。preds.cpu().numpy()把預(yù)測結(jié)果搬到CPU并轉(zhuǎn)成numpy這是傳給sklearn做指標(biāo)計(jì)算的標(biāo)準(zhǔn)姿勢。分類報告會輸出每個類別的精確率、召回率和F1混淆矩陣能直接看出哪些類別互相混淆——比如雛菊和蒲公英經(jīng)常被搞混那就該去檢查這兩個類別的樣本是不是拍攝角度太接近。5. AlexNet花卉分類避坑記錄4個常見的訓(xùn)練與推理問題5.1 驗(yàn)證集準(zhǔn)確率卡住不動學(xué)習(xí)率與歸一化的鍋現(xiàn)象訓(xùn)練Loss在下降但驗(yàn)證集準(zhǔn)確率從第1個epoch到第10個epoch都紋絲不動穩(wěn)定在某個低值附近。原因常見的有兩個。一是預(yù)處理里漏了Normalize原始RGB像素值范圍是0到255預(yù)訓(xùn)練權(quán)重期望的是標(biāo)準(zhǔn)化后的分布輸入差異太大會讓模型一直找不到有效梯度方向二是學(xué)習(xí)率對當(dāng)前數(shù)據(jù)量來說偏大損失在loss曲面里震蕩驗(yàn)證集指標(biāo)跟著來回抖。解決先檢查transform里有沒有帶上ImageNet的均值和標(biāo)準(zhǔn)差沒有就補(bǔ)上。學(xué)習(xí)率從0.001降到0.0005或0.0003并確認(rèn)scheduler確實(shí)在step。還有一個容易忽略的點(diǎn)驗(yàn)證時也要走val_transform不能拿訓(xùn)練用的RandomResizedCrop去驗(yàn)證隨機(jī)裁剪會讓驗(yàn)證指標(biāo)每次都變。5.2 訓(xùn)練集99%驗(yàn)證集65%數(shù)據(jù)量撐不起全參數(shù)微調(diào)現(xiàn)象訓(xùn)練集準(zhǔn)確率一路漲到99%以上驗(yàn)證集卻停留在60%多兩個曲線之間的gap越來越大。這是最典型的過擬合長相。原因AlexNet全模型有近6000萬參數(shù)如果每類只有50張圖模型完全有能力把訓(xùn)練圖背下來但學(xué)不到能泛化的花形特征。數(shù)據(jù)量越小過擬合來得越快。解決按優(yōu)先級排序做三件事。先把訓(xùn)練集數(shù)據(jù)增強(qiáng)加滿RandomResizedCrop的scale收緊到(0.5, 1.0)加上翻轉(zhuǎn)和顏色抖動。然后把features部分凍結(jié)只訓(xùn)練最后的分類器用for param in model.features.parameters(): param.requires_grad False這樣實(shí)際參與訓(xùn)練的參數(shù)只有百萬級。最后把第一個Dropout從0.5提高到0.7。這三步做完驗(yàn)證集通常能拉回85%以上。5.3 加載預(yù)訓(xùn)練權(quán)重報尺寸不匹配96核與64核的歷史坑現(xiàn)象自己按論文結(jié)構(gòu)寫了nn.Conv2d(3, 96, kernel_size11, stride4)然后執(zhí)行model.load_state_dict(torch.load(alexnet_weights.pth))報錯提示features.0.weight尺寸不匹配期望[96,3,11,11]實(shí)際是[64,3,11,11]。原因原創(chuàng)論文因?yàn)楫?dāng)年GPU顯存限制把第一層96個卷積核拆到兩張顯卡上torchvision實(shí)現(xiàn)的是單卡版本第一層只有64個核。網(wǎng)上流傳的不少代碼結(jié)構(gòu)抄的是論文原版而權(quán)重文件來自torchvision兩者天然對不上。解決直接用torchvision.models.alexnet加載官方結(jié)構(gòu)不要自己寫features部分。如果堅(jiān)持用論文96核結(jié)構(gòu)就只能從隨機(jī)初始化開始訓(xùn)放棄預(yù)訓(xùn)練帶來的遷移收益。在花卉這種小數(shù)據(jù)集上放棄預(yù)訓(xùn)練權(quán)重基本等于放棄最好的提升路徑。5.4 推理結(jié)果全亂預(yù)處理不一致與漏掉model.eval()現(xiàn)象訓(xùn)練時驗(yàn)證集準(zhǔn)確率90%以上把圖片單獨(dú)拎出來走一遍推理代碼結(jié)果不是錯得離譜就是每次都預(yù)測同一個類別。原因兩個很隱蔽的錯。一是推理代碼用的預(yù)處理和訓(xùn)練時不一致比如訓(xùn)練用了Resize(256)CenterCrop(227)推理時直接Resize((227,227))圖片比例被壓扁變形二是model.eval()沒調(diào)用Dropout在推理時還在隨機(jī)丟棄神經(jīng)元輸出自然不穩(wěn)定。解決把推理用的transform參數(shù)和驗(yàn)證集對齊嚴(yán)格保持Resize((256,256)) → CenterCrop(227) → Normalize的順序。調(diào)用模型前先model.eval()并且推理全程包在torch.no_grad()里。還有一個細(xì)節(jié)用PIL打開圖片后記得.convert(RGB)否則遇到帶透明通道的PNG或灰度圖通道數(shù)不一致會直接報錯。6. 收尾技巧把訓(xùn)練好的權(quán)重封裝成單圖分類接口6.1 保存類別名權(quán)重和類別順序必須一起落盤訓(xùn)練腳本里很多人只存了model.state_dict()沒存類別名。權(quán)重文件里只有數(shù)字索引0對應(yīng)哪個花名完全取決于訓(xùn)練時的ImageFolder目錄順序。換個環(huán)境重跑數(shù)據(jù)加載類別順序一變化索引全錯位模型表現(xiàn)像完全沒訓(xùn)練過一樣。我的習(xí)慣是訓(xùn)練結(jié)束順手把類別名落盤import json with open(class_names.json, w, encodingutf-8) as f: json.dump(train_ds.classes, f, ensure_asciiFalse, indent2)train_ds.classes是ImageFolder按文件夾名排序后的類別列表和模型輸出的索引一一對應(yīng)。把這個文件放到和權(quán)重同一個目錄推理時只認(rèn)文件里的順序不依賴環(huán)境里重新加載數(shù)據(jù)集。6.2 推理封裝從單張圖片路徑到Top-K結(jié)果封裝一個獨(dú)立的分類器類把權(quán)重加載、預(yù)處理、推理合并在一起后續(xù)不管是做Web接口還是批量跑目錄圖片都復(fù)用同一個類import json import torch import torch.nn as nn from PIL import Image from torchvision import models, transforms class FlowerClassifier: def __init__(self, weights_path, class_names_path, deviceNone): self.device device or torch.device( cuda if torch.cuda.is_available() else cpu) with open(class_names_path, r, encodingutf-8) as f: self.class_names json.load(f) self.model models.alexnet(weightsNone) in_features self.model.classifier[6].in_features self.model.classifier[6] nn.Linear(in_features, len(self.class_names)) self.model.load_state_dict( torch.load(weights_path, map_locationself.device)) self.model.to(self.device) self.model.eval() self.transform transforms.Compose([ transforms.Resize((256, 256)), transforms.CenterCrop(227), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) def predict_topk(self, image_path, top_k3): image Image.open(image_path).convert(RGB) x self.transform(image).unsqueeze(0).to(self.device) with torch.no_grad(): logits self.model(x) probs torch.softmax(logits, dim1) top_probs, top_indices torch.topk(probs, top_k, dim1) return [(self.class_names[idx], float(prob)) for prob, idx in zip(top_probs[0], top_indices[0])]調(diào)用方式很直接classifier FlowerClassifier(best_alexnet_flower.pth, class_names.json) for name, prob in classifier.predict_topk(test_rose.jpg, top_k3): print(f{name}: {prob:.2%})參數(shù)說明map_locationself.device讓權(quán)重文件在只有CPU的機(jī)器上也能加載。unsqueeze(0)給單張圖片補(bǔ)上batch維度因?yàn)槟Ｐ鸵筝斎胧撬木S張量。top_k3返回概率最高的三個類別在演示場景里比只給一個答案更有說服力——如果排名第一的置信度不足60%排名第二的往往也有參考價值。這套方案做到今天給我感觸最深的不是網(wǎng)絡(luò)結(jié)構(gòu)本身而是數(shù)據(jù)管道和評估閉環(huán)一定要做扎實(shí)。我最早跑通時偷懶沒存類別名換機(jī)器推理時輸出全亂排查了一下午才發(fā)現(xiàn)是索引錯位從那以后權(quán)重和類別名永遠(yuǎn)一起落盤。希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取