簡介基于AlexNet的花卉分類識別系統(tǒng)是一套完整的預(yù)訓(xùn)練模型應(yīng)用項目面向Python與深度學(xué)習(xí)入門者、進(jìn)階學(xué)習(xí)者以及正在籌備畢業(yè)設(shè)計或課程設(shè)計的同學(xué)可快速用于圖像分類實踐或項目立項。系統(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ù)自帶樣例快速驗證效果。已有346人學(xué)習(xí)適合用作入門深度學(xué)習(xí)圖像分類的基礎(chǔ)參考也可作為課程設(shè)計或工程實訓(xù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é)生另一類是想快速驗證花卉識別效果、不打算從零造輪子的開發(fā)者。下面按數(shù)據(jù)準(zhǔn)備、網(wǎng)絡(luò)實現(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在普通消費級顯卡上十幾分鐘就夠適合先把整條鏈路跑通。102類的版本留到后面驗證模型泛化能力時再用。如果決定自建數(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)練集和驗證集數(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)}張圖建議補充) 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圖片的路徑收集起來隨機打亂后按8:2比例分別復(fù)制到train和val下對應(yīng)的類別文件夾。用copy2而不是move是為了保留一份原始素材之后補充了新圖想重新切分時不需要重新采集。參數(shù)說明split 0.8是經(jīng)驗值。每類樣本不足50張時我會改成0.85盡量多留訓(xùn)練樣本樣本充足時0.7也可以驗證集更有統(tǒng)計意義。腳本里漏掉格式后綴是常見的翻車點如果原始圖里有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等于每輪從原圖隨機位置、隨機大小裁一塊區(qū)域是一種不需要額外數(shù)據(jù)就能生效的數(shù)據(jù)增強驗證時不希望結(jié)果帶隨機性所以固定用CenterCrop裁正中間。Normalize的均值和標(biāo)準(zhǔn)差取自ImageNet全量統(tǒng)計只要用預(yù)訓(xùn)練權(quán)重這組值就不能動改了等于把輸入分布整體偏移訓(xùn)練會明顯變慢。參數(shù)說明RandomResizedCrop默認(rèn)scale范圍是(0.08, 1.0)。對花朵這類主體通常占畫面較大面積的場景我會把scale收緊到(0.5, 1.0)避免裁出來的小塊里幾乎沒有花。ColorJitter里saturation建議克制一點花卉識別對顏色敏感飽和度抖動過猛會讓模型學(xué)到失真的顏色分布。3. AlexNet模型搭建用PyTorch逐層實現(xiàn)卷積與全連接3.1 AlexNet結(jié)構(gòu)拆解從卷積到全連接的層與尺寸AlexNet的主體可以拆成特征提取器加分類器兩段。特征提取器是5層卷積前兩層后面接MaxPool第3、4層卷積不接池化第5層卷積后面再接一次MaxPool分類器是3層全連接前兩層各接一個ReLU和一個Dropout。整個網(wǎng)絡(luò)的核心思想很簡單用卷積逐步把空間尺寸壓小、把通道數(shù)加多最后展平成全連接向量做分類。網(wǎng)上不少實現(xiàn)把第一層卷積核數(shù)量寫成96這是照抄原始論文的結(jié)果。原始論文因為當(dāng)年單卡顯存不夠把96個卷積核拆到兩張GPU上并行訓(xùn)練torchvision的現(xiàn)代單卡實現(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逐層實現(xiàn)可直接復(fù)制的AlexNet代碼實際項目里我不會自己從零寫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è)計——卷積層本身有空間冗余對Dropout不敏感全連接層參數(shù)占比最大、最容易過擬合所以只在分類器里加。參數(shù)說明dropout0.5是原始論文的值在遷移學(xué)習(xí)場景下如果數(shù)據(jù)量不大我會把前面那個Dropout提到0.7后面那個保持0.5。原始論文里還有一個Local Response Normalization層現(xiàn)代框架里它對最終精度的貢獻(xiàn)很小還會拖慢訓(xùn)練和推理速度我一般直接省略實測對花卉分類這類小類別任務(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ù)訓(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張再放開全模型微調(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)鍵是把驗證邏輯獨立成函數(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做一次清零、前向、算損失、反向、更新五步。驗證和保存都放在epoch層面做以驗證集準(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一路下降驗證Loss先降后平緩驗證準(zhǔn)確率同步爬升。如果訓(xùn)練Loss還在降、驗證Loss卻掉頭向上說明模型開始背訓(xùn)練集了這時候繼續(xù)訓(xùn)練只會讓驗證集越來越差。Early Stopping的實現(xiàn)不復(fù)雜也不需要額外裝庫記錄當(dāng)前最佳驗證準(zhǔn)確率連續(xù)N個epoch沒有刷新就提前終止。N我一般取7到10。如果你用的是上面那個腳本torch.save已經(jīng)天然只保留最佳權(quán)重所以Early Stopping只是幫你省訓(xùn)練時間不省也可以但數(shù)據(jù)量小的時候多跑那幾個epoch純屬浪費電。4.3 評估不止看準(zhǔn)確率混淆矩陣與單類召回率準(zhǔn)確率是平均數(shù)最容易騙人。如果數(shù)據(jù)集里玫瑰占了一半模型把所有圖片都判成玫瑰準(zhǔn)確率也有50%但你真正關(guān)心的可能是郁金香一張都認(rèn)不出來?；ɑ芊诸惖脑u估我至少要打印一份分類報告和混淆矩陣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)計算的標(biāo)準(zhǔn)姿勢。分類報告會輸出每個類別的精確率、召回率和F1混淆矩陣能直接看出哪些類別互相混淆——比如雛菊和蒲公英經(jīng)常被搞混那就該去檢查這兩個類別的樣本是不是拍攝角度太接近。5. AlexNet花卉分類避坑記錄4個常見的訓(xùn)練與推理問題5.1 驗證集準(zhǔn)確率卡住不動學(xué)習(xí)率與歸一化的鍋現(xiàn)象訓(xùn)練Loss在下降但驗證集準(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曲面里震蕩驗證集指標(biāo)跟著來回抖。解決先檢查transform里有沒有帶上ImageNet的均值和標(biāo)準(zhǔn)差沒有就補上。學(xué)習(xí)率從0.001降到0.0005或0.0003并確認(rèn)scheduler確實在step。還有一個容易忽略的點驗證時也要走val_transform不能拿訓(xùn)練用的RandomResizedCrop去驗證隨機裁剪會讓驗證指標(biāo)每次都變。5.2 訓(xùn)練集99%驗證集65%數(shù)據(jù)量撐不起全參數(shù)微調(diào)現(xiàn)象訓(xùn)練集準(zhǔn)確率一路漲到99%以上驗證集卻停留在60%多兩個曲線之間的gap越來越大。這是最典型的過擬合長相。原因AlexNet全模型有近6000萬參數(shù)如果每類只有50張圖模型完全有能力把訓(xùn)練圖背下來但學(xué)不到能泛化的花形特征。數(shù)據(jù)量越小過擬合來得越快。解決按優(yōu)先級排序做三件事。先把訓(xùn)練集數(shù)據(jù)增強加滿RandomResizedCrop的scale收緊到(0.5, 1.0)加上翻轉(zhuǎn)和顏色抖動。然后把features部分凍結(jié)只訓(xùn)練最后的分類器用for param in model.features.parameters(): param.requires_grad False這樣實際參與訓(xùn)練的參數(shù)只有百萬級。最后把第一個Dropout從0.5提高到0.7。這三步做完驗證集通常能拉回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]實際是[64,3,11,11]。原因原創(chuàng)論文因為當(dāng)年GPU顯存限制把第一層96個卷積核拆到兩張顯卡上torchvision實現(xiàn)的是單卡版本第一層只有64個核。網(wǎng)上流傳的不少代碼結(jié)構(gòu)抄的是論文原版而權(quán)重文件來自torchvision兩者天然對不上。解決直接用torchvision.models.alexnet加載官方結(jié)構(gòu)不要自己寫features部分。如果堅持用論文96核結(jié)構(gòu)就只能從隨機初始化開始訓(xùn)放棄預(yù)訓(xùn)練帶來的遷移收益。在花卉這種小數(shù)據(jù)集上放棄預(yù)訓(xùn)練權(quán)重基本等于放棄最好的提升路徑。5.4 推理結(jié)果全亂預(yù)處理不一致與漏掉model.eval()現(xiàn)象訓(xùn)練時驗證集準(zhǔn)確率90%以上把圖片單獨拎出來走一遍推理代碼結(jié)果不是錯得離譜就是每次都預(yù)測同一個類別。原因兩個很隱蔽的錯。一是推理代碼用的預(yù)處理和訓(xùn)練時不一致比如訓(xùn)練用了Resize(256)CenterCrop(227)推理時直接Resize((227,227))圖片比例被壓扁變形二是model.eval()沒調(diào)用Dropout在推理時還在隨機丟棄神經(jīng)元輸出自然不穩(wěn)定。解決把推理用的transform參數(shù)和驗證集對齊嚴(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é)果封裝一個獨立的分類器類把權(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的機器上也能加載。unsqueeze(0)給單張圖片補上batch維度因為模型要求輸入是四維張量。top_k3返回概率最高的三個類別在演示場景里比只給一個答案更有說服力——如果排名第一的置信度不足60%排名第二的往往也有參考價值。這套方案做到今天給我感觸最深的不是網(wǎng)絡(luò)結(jié)構(gòu)本身而是數(shù)據(jù)管道和評估閉環(huán)一定要做扎實。我最早跑通時偷懶沒存類別名換機器推理時輸出全亂排查了一下午才發(fā)現(xiàn)是索引錯位從那以后權(quán)重和類別名永遠(yuǎn)一起落盤。希望幫到你。本文還有配套的精品資源點擊獲取