簡(jiǎn)介這是一套面向高校學(xué)生與初學(xué)者的垃圾分類深度學(xué)習(xí)實(shí)戰(zhàn)項(xiàng)目源碼基于Python與主流深度學(xué)習(xí)框架實(shí)現(xiàn)可直接用于畢業(yè)設(shè)計(jì)、期末大作業(yè)或課程設(shè)計(jì)場(chǎng)景。項(xiàng)目已通過教師指導(dǎo)與驗(yàn)收屬于高分完整方案對(duì)零基礎(chǔ)讀者也較為友好能幫助快速理解圖像分類任務(wù)的完整落地流程。壓縮包共收錄6319個(gè)文件約18.98MB其中以py源碼與pyc編譯文件為主體輔以mo、po多語言資源、html與js前端頁面、png與svg圖像素材以及txt說明、css樣式和少量exe工具整體結(jié)構(gòu)接近可直接運(yùn)行的工程形態(tài)。目前已有338人學(xué)習(xí)下載說明該方案在同類選題中具備一定參考熱度。讀者可從中獲取完整的垃圾分類識(shí)別代碼、模型訓(xùn)練與推理腳本、界面展示頁面及配套資源文件便于對(duì)照復(fù)現(xiàn)實(shí)驗(yàn)、梳理項(xiàng)目目錄組織方式并在此基礎(chǔ)上進(jìn)行功能擴(kuò)展或二次開發(fā)。1. 從一份能跑通的垃圾分類源碼說起它到底解決了什么如果你正在為畢業(yè)設(shè)計(jì)發(fā)愁或者想找一個(gè)能真正跑起來、有完整訓(xùn)練流程的深度學(xué)習(xí)項(xiàng)目練手這份基于 Python 與深度學(xué)習(xí)的垃圾分類系統(tǒng)源碼值得先看一眼。它不是那種只丟幾個(gè).py文件、連數(shù)據(jù)集路徑都寫死的半成品而是一套從數(shù)據(jù)預(yù)處理、模型訓(xùn)練、評(píng)估到推理預(yù)測(cè)都能串起來的完整工程。垃圾分類本身是個(gè)典型的圖像分類任務(wù)類別通常包括可回收物、廚余垃圾、有害垃圾和其他垃圾四類正好覆蓋了 CNN 從入門到落地的全部關(guān)鍵環(huán)節(jié)。你拿到手之后最直接的用途就是改一改數(shù)據(jù)集路徑、調(diào)一調(diào)超參數(shù)就能跑出自己的分類模型往深了說這套代碼的結(jié)構(gòu)足夠清晰方便你在上面加注意力機(jī)制、換 backbone、做消融實(shí)驗(yàn)支撐起一篇像樣的畢業(yè)論文。適合誰適合已經(jīng)裝好 Python 環(huán)境、知道pip install怎么用、但還沒完整走過一遍深度學(xué)習(xí)項(xiàng)目流程的人。下面我就按實(shí)際拆包復(fù)現(xiàn)的順序把這份資源從里到外講一遍。2. 拆開壓縮包先看什么目錄結(jié)構(gòu)與技術(shù)棧選型2.1 拿到源碼后的第一輪文件排查解壓之后別急著python train.py先花五分鐘把目錄結(jié)構(gòu)過一遍。這類畢業(yè)設(shè)計(jì)項(xiàng)目通常包含以下幾個(gè)核心部分data/或dataset/放原始圖片按類別分文件夾models/存網(wǎng)絡(luò)定義utils/放數(shù)據(jù)加載和預(yù)處理腳本根目錄下有train.py、predict.py、evaluate.py這幾個(gè)入口文件可能還有一個(gè)config.py或args.py統(tǒng)一管理超參數(shù)。我一般會(huì)先執(zhí)行下面這段命令把文件樹和文件大小摸清楚# 查看目錄結(jié)構(gòu)排除緩存文件 find . -type f -not -path ./.git/* -not -name *.pyc | head -50 # 統(tǒng)計(jì)各類圖片數(shù)量確認(rèn)數(shù)據(jù)集是否完整 for dir in data/*/; do echo -n $dir: find $dir -type f \( -name *.jpg -o -name *.png -o -name *.jpeg \) | wc -l done第一段命令幫你快速定位入口腳本和配置文件的位置第二段命令按類別統(tǒng)計(jì)圖片數(shù)量。如果某個(gè)類別只有幾十張圖那訓(xùn)練時(shí)大概率會(huì)嚴(yán)重過擬合需要先做數(shù)據(jù)增強(qiáng)或者補(bǔ)充樣本。這一步看著簡(jiǎn)單但我見過太多人直接開訓(xùn)跑了半天才發(fā)現(xiàn)某個(gè)類別文件夾是空的血淚經(jīng)驗(yàn)。2.2 技術(shù)棧選型為什么是 PyTorch CNN這份源碼大概率用的是 PyTorch 而不是 TensorFlow原因很實(shí)際PyTorch 的動(dòng)態(tài)圖機(jī)制對(duì)調(diào)試更友好print一下 tensor 的 shape 就能定位問題而靜態(tài)圖框架報(bào)錯(cuò)信息往往讓人摸不著頭腦。網(wǎng)絡(luò)結(jié)構(gòu)方面垃圾分類這種四分類任務(wù)用 ResNet18 或 MobileNetV2 做遷移學(xué)習(xí)是最穩(wěn)妥的選擇。ResNet18 參數(shù)量約 1100 萬在 ImageNet 上預(yù)訓(xùn)練過的權(quán)重拿來微調(diào)通常 20 到 30 個(gè) epoch 就能收斂到一個(gè)不錯(cuò)的精度。MobileNetV2 更輕量適合你想部署到邊緣設(shè)備或者對(duì)推理速度有要求的場(chǎng)景。選型理由可以歸結(jié)為三點(diǎn)第一數(shù)據(jù)量通常不大從零訓(xùn)練一個(gè)深層網(wǎng)絡(luò)必然過擬合遷移學(xué)習(xí)是正解第二垃圾分類的類間差異比較明顯不需要特別復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu)第三ResNet 的殘差連接能有效緩解梯度消失訓(xùn)練過程更穩(wěn)定。如果你在論文里需要對(duì)比實(shí)驗(yàn)可以再加一個(gè) VGG16 或者自己搭一個(gè)簡(jiǎn)單的四層 CNN 作為 baseline這樣論文的實(shí)驗(yàn)部分會(huì)更充實(shí)。2.3 環(huán)境配置與依賴安裝的實(shí)操步驟環(huán)境配置是新手翻車最多的地方。我建議用 conda 建一個(gè)獨(dú)立環(huán)境避免和系統(tǒng) Python 打架# 創(chuàng)建虛擬環(huán)境指定 Python 3.8 或 3.9 conda create -n garbage_cls python3.9 -y conda activate garbage_cls # 安裝 PyTorch根據(jù)你的 CUDA 版本選擇對(duì)應(yīng)命令 # 如果沒有 GPU用 CPU 版本即可 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 安裝其他依賴 pip install numpy pandas matplotlib opencv-python pillow tqdm scikit-learn這里有幾個(gè)參數(shù)需要留意python3.9是我實(shí)測(cè)兼容性最好的版本3.10 以上有時(shí)會(huì)遇到某些庫還沒適配的問題PyTorch 的安裝命令要根據(jù)你的顯卡驅(qū)動(dòng)來選cu118對(duì)應(yīng) CUDA 11.8如果你不確定直接去 PyTorch 官網(wǎng)用它的選擇器生成命令最保險(xiǎn)。安裝完之后跑一句python -c import torch; print(torch.cuda.is_available())返回True說明 GPU 可用返回False就是 CPU 模式訓(xùn)練速度會(huì)慢很多但代碼本身不用改。提示如果你用的是 VSCode記得在右下角切換 Python 解釋器到剛創(chuàng)建的 conda 環(huán)境否則終端里裝好了、編輯器里還是找不到包。3. 數(shù)據(jù)管道與模型訓(xùn)練從圖片到可用的分類器3.1 數(shù)據(jù)預(yù)處理與增強(qiáng)策略垃圾分類數(shù)據(jù)集通常存在兩個(gè)問題類別不平衡和樣本量不足。源碼里一般會(huì)用torchvision.transforms做在線增強(qiáng)常見組合是隨機(jī)裁剪、水平翻轉(zhuǎn)、顏色抖動(dòng)。我一般會(huì)在這個(gè)基礎(chǔ)上再加一個(gè)隨機(jī)旋轉(zhuǎn)因?yàn)槔鴪D片的拍攝角度往往很隨意。下面是一段典型的數(shù)據(jù)加載代碼import torch from torchvision import transforms, datasets from torch.utils.data import DataLoader, WeightedRandomSampler # 訓(xùn)練集增強(qiáng)隨機(jī)裁剪 翻轉(zhuǎn) 旋轉(zhuǎn) 顏色抖動(dòng) train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), # 隨機(jī)裁剪并縮放到224 transforms.RandomHorizontalFlip(p0.5), # 50%概率水平翻轉(zhuǎn) transforms.RandomRotation(15), # 隨機(jī)旋轉(zhuǎn)±15度 transforms.ColorJitter(brightness0.2, contrast0.2), # 亮度對(duì)比度擾動(dòng) transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], # ImageNet標(biāo)準(zhǔn)均值 std[0.229, 0.224, 0.225]) # ImageNet標(biāo)準(zhǔn)方差 ]) # 驗(yàn)證集只做縮放和歸一化不做增強(qiáng) 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]) ]) # 加載數(shù)據(jù)集ImageFolder要求按類別分文件夾 train_dataset datasets.ImageFolder(data/train, transformtrain_transform) val_dataset datasets.ImageFolder(data/val, transformval_transform) # 處理類別不平衡給少數(shù)類更高采樣權(quán)重 targets [s[1] for s in train_dataset.samples] class_counts torch.bincount(torch.tensor(targets)) weights 1.0 / class_counts.float() sample_weights weights[targets] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader DataLoader(train_dataset, batch_size32, samplersampler, num_workers4) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4)這段代碼的關(guān)鍵點(diǎn)有三個(gè)RandomResizedCrop的scale參數(shù)控制裁剪區(qū)域占原圖的比例設(shè)太小會(huì)把垃圾物體裁掉設(shè)太大增強(qiáng)效果不明顯0.7 到 1.0 是比較安全的范圍Normalize用的均值和方差必須和預(yù)訓(xùn)練模型一致否則遷移學(xué)習(xí)的效果會(huì)打折扣WeightedRandomSampler是解決類別不平衡的利器它讓少數(shù)類樣本在每批中被抽到的概率更高比簡(jiǎn)單的過采樣更不容易過擬合。num_workers設(shè)成 4 還是 8 取決于你機(jī)器的 CPU 核心數(shù)設(shè)太大反而會(huì)因?yàn)檫M(jìn)程切換拖慢速度。3.2 模型構(gòu)建與遷移學(xué)習(xí)微調(diào)模型部分的核心思路是加載預(yù)訓(xùn)練權(quán)重替換最后的全連接層然后分階段微調(diào)。源碼里可能直接寫死了 ResNet18但你可以根據(jù)自己需求換成 ResNet50 或 EfficientNet。下面是我常用的寫法import torch.nn as nn from torchvision import models def build_model(num_classes4, backboneresnet18, pretrainedTrue): if backbone resnet18: model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None) in_features model.fc.in_features # 替換最后的全連接層輸出類別數(shù)改為4 model.fc nn.Sequential( nn.Dropout(0.3), # 防止過擬合 nn.Linear(in_features, num_classes) ) elif backbone mobilenet_v2: model models.mobilenet_v2(weightsmodels.MobileNet_V2_Weights.IMAGENET1K_V1 if pretrained else None) in_features model.classifier[1].in_features model.classifier nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, num_classes) ) return model # 分階段微調(diào)先凍結(jié)骨干網(wǎng)絡(luò)只訓(xùn)練分類頭 model build_model(num_classes4, backboneresnet18) for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True # 第一階段優(yōu)化器只更新分類頭參數(shù) optimizer torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3)這里的分階段策略值得展開說第一階段凍結(jié)骨干網(wǎng)絡(luò)只訓(xùn)練新加的全連接層學(xué)習(xí)率可以設(shè)大一點(diǎn)1e-3因?yàn)殡S機(jī)初始化的分類頭需要快速收斂訓(xùn)練 5 個(gè) epoch 之后解凍骨干網(wǎng)絡(luò)的后幾層用更小的學(xué)習(xí)率1e-4 或 1e-5做微調(diào)避免把預(yù)訓(xùn)練學(xué)到的通用特征破壞掉。Dropout(0.3)是防止過擬合的常規(guī)操作如果你的數(shù)據(jù)集特別小每類不到 200 張可以加到 0.5。損失函數(shù)用CrossEntropyLoss就夠了如果類別不平衡特別嚴(yán)重可以給它傳weight參數(shù)和上面的采樣器二選一即可不用同時(shí)上。3.3 訓(xùn)練循環(huán)與關(guān)鍵參數(shù)設(shè)置訓(xùn)練循環(huán)的代碼看起來都差不多但魔鬼在細(xì)節(jié)里。下面這段是我會(huì)實(shí)際用的版本import torch from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30) best_acc 0.0 for epoch in range(30): model.train() running_loss 0.0 for images, labels in tqdm(train_loader, descfEpoch {epoch1}): 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() # 驗(yàn)證階段 model.eval() correct, total 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total scheduler.step() # 保存最佳模型 if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fBest model saved with acc: {best_acc:.4f})幾個(gè)參數(shù)需要根據(jù)實(shí)際情況調(diào)整batch_size設(shè) 32 是 8GB 顯存下的安全值顯存夠大可以上 64學(xué)習(xí)率調(diào)度器用余弦退火比固定學(xué)習(xí)率效果好T_max設(shè)成總 epoch 數(shù)保存模型時(shí)只存state_dict而不是整個(gè)模型這樣加載時(shí)更靈活。如果你發(fā)現(xiàn)訓(xùn)練 loss 一直在降但驗(yàn)證 acc 不漲大概率是過擬合了可以加數(shù)據(jù)增強(qiáng)、加 Dropout、或者提前停止。反過來如果訓(xùn)練 loss 都不降先檢查學(xué)習(xí)率是不是太大或者數(shù)據(jù)標(biāo)簽有沒有搞錯(cuò)。4. 推理部署與效果驗(yàn)證模型訓(xùn)完之后怎么用4.1 單張圖片推理與批量預(yù)測(cè)訓(xùn)練完拿到best_model.pth之后下一步就是寫推理腳本。很多人訓(xùn)完模型就不知道怎么辦了其實(shí)推理代碼比訓(xùn)練簡(jiǎn)單得多from PIL import Image import torch from torchvision import transforms # 加載模型結(jié)構(gòu)再加載權(quán)重 model build_model(num_classes4, backboneresnet18, pretrainedFalse) model.load_state_dict(torch.load(best_model.pth, map_locationcpu)) model.eval() # 推理預(yù)處理和驗(yàn)證集保持一致 infer_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]) ]) class_names [可回收物, 廚余垃圾, 有害垃圾, 其他垃圾] def predict(image_path): image Image.open(image_path).convert(RGB) tensor infer_transform(image).unsqueeze(0) # 增加batch維度 with torch.no_grad(): outputs model(tensor) probabilities torch.softmax(outputs, dim1) confidence, predicted torch.max(probabilities, 1) return class_names[predicted.item()], confidence.item() # 測(cè)試 label, conf predict(test.jpg) print(f預(yù)測(cè)類別: {label}, 置信度: {conf:.4f})這里有個(gè)容易忽略的點(diǎn)load_state_dict之前必須先實(shí)例化模型結(jié)構(gòu)而且pretrained要設(shè)成False否則會(huì)去下載預(yù)訓(xùn)練權(quán)重浪費(fèi)時(shí)間。map_locationcpu是防止在沒有 GPU 的機(jī)器上加載時(shí)報(bào)錯(cuò)。推理時(shí)的預(yù)處理必須和驗(yàn)證集完全一致包括Resize的大小和Normalize的參數(shù)差一點(diǎn)都會(huì)導(dǎo)致精度下降。torch.softmax把輸出轉(zhuǎn)成概率分布置信度低于 0.6 的時(shí)候我一般會(huì)提示“結(jié)果不確定”這在演示系統(tǒng)里很實(shí)用。4.2 評(píng)估指標(biāo)與混淆矩陣只看準(zhǔn)確率是不夠的尤其是類別不平衡的時(shí)候。我習(xí)慣在驗(yàn)證集上跑一遍完整的評(píng)估輸出分類報(bào)告和混淆矩陣from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) _, predicted torch.max(outputs, 1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.numpy()) # 打印每個(gè)類別的精確率、召回率、F1 print(classification_report(all_labels, all_preds, target_namesclass_names)) # 畫混淆矩陣 cm confusion_matrix(all_labels, all_preds) sns.heatmap(cm, annotTrue, fmtd, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(預(yù)測(cè)) plt.ylabel(真實(shí)) plt.savefig(confusion_matrix.png, dpi150, bbox_inchestight)classification_report會(huì)給出每個(gè)類別的 precision、recall 和 f1-score如果某個(gè)類別的 recall 特別低說明模型把很多該類樣本漏判了可能是樣本太少或者特征不明顯?；煜仃嚹苤庇^看出模型把哪兩類搞混了比如“廚余垃圾”和“其他垃圾”經(jīng)?；ハ嗾`判這時(shí)候可以考慮針對(duì)這兩類補(bǔ)充訓(xùn)練數(shù)據(jù)或者加一個(gè)二分類器做后處理。這些分析寫進(jìn)論文里比單純報(bào)一個(gè)準(zhǔn)確率有說服力得多。4.3 用 Grad-CAM 做可視化解釋答辯的時(shí)候老師經(jīng)常會(huì)問“模型到底學(xué)到了什么”這時(shí)候 Grad-CAM 熱力圖就是你的后悔藥。它能把模型關(guān)注區(qū)域高亮出來讓你看到模型是根據(jù)哪些像素做判斷的from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.image import show_cam_on_image import numpy as np import cv2 # 選擇目標(biāo)層ResNet18的最后一個(gè)卷積層 target_layers [model.layer4[-1]] cam GradCAM(modelmodel, target_layerstarget_layers) # 準(zhǔn)備輸入 rgb_img cv2.imread(test.jpg) rgb_img cv2.resize(rgb_img, (224, 224)) input_tensor infer_transform(Image.fromarray(cv2.cvtColor(rgb_img, cv2.COLOR_BGR2RGB))).unsqueeze(0) # 生成熱力圖 grayscale_cam cam(input_tensorinput_tensor)[0, :] visualization show_cam_on_image(rgb_img / 255.0, grayscale_cam, use_rgbTrue) cv2.imwrite(gradcam_result.jpg, cv2.cvtColor(visualization, cv2.COLOR_RGB2BGR))target_layers選最后一層卷積是因?yàn)樗Ａ袅俗疃嗟目臻g信息選太淺的層熱力圖會(huì)很散。生成的熱力圖疊加在原圖上紅色區(qū)域就是模型最關(guān)注的地方。如果模型關(guān)注的是垃圾物體本身說明學(xué)對(duì)了如果關(guān)注的是背景或者水印那就要檢查數(shù)據(jù)集里有沒有引入偏差。這個(gè)工具在論文里放一張圖比寫一堆文字解釋都管用。5. 避坑與常見問題排查5.1 訓(xùn)練 loss 不下降或震蕩嚴(yán)重現(xiàn)象訓(xùn)練幾個(gè) epoch 后 loss 一直在 2.0 附近徘徊或者上下劇烈跳動(dòng)。原因通常是學(xué)習(xí)率設(shè)太大了或者數(shù)據(jù)標(biāo)簽有問題。先檢查學(xué)習(xí)率遷移學(xué)習(xí)第一階段用 1e-3 是上限如果 loss 炸了降到 1e-4 試試。然后抽查一批數(shù)據(jù)把圖片和標(biāo)簽打印出來看看是否對(duì)應(yīng)我遇到過有人把類別文件夾名字寫錯(cuò)導(dǎo)致標(biāo)簽全亂的。解決用torch.optim.lr_scheduler.ReduceLROnPlateau做自適應(yīng)降學(xué)習(xí)率同時(shí)寫一個(gè)check_dataset.py腳本隨機(jī)可視化 16 張圖確認(rèn)標(biāo)簽。5.2 驗(yàn)證集準(zhǔn)確率遠(yuǎn)低于訓(xùn)練集現(xiàn)象訓(xùn)練集準(zhǔn)確率 95%驗(yàn)證集只有 60%。這是典型的過擬合。原因可能是數(shù)據(jù)增強(qiáng)不夠、模型參數(shù)量太大、或者訓(xùn)練集和驗(yàn)證集分布不一致。解決先加數(shù)據(jù)增強(qiáng)RandomResizedCrop、ColorJitter、RandomRotation都加上然后加 Dropout 和權(quán)重衰減weight_decay1e-4如果還不行換更小的模型比如從 ResNet50 換回 ResNet18。另外檢查一下驗(yàn)證集的圖片是不是從訓(xùn)練集里漏過來的雖然聽起來離譜但我確實(shí)見過。5.3 GPU 顯存不足報(bào) CUDA out of memory現(xiàn)象訓(xùn)練到一半突然報(bào)RuntimeError: CUDA out of memory。原因可能是 batch_size 太大或者沒有釋放中間變量。解決先把batch_size減半從 32 降到 16然后在訓(xùn)練循環(huán)里加torch.cuda.empty_cache()雖然它不能完全解決問題但能緩解碎片化如果還不行用torch.cuda.amp做混合精度訓(xùn)練顯存占用能降差不多一半。代碼改動(dòng)很小在loss.backward()前加一個(gè)scaler.scale(loss).backward()就行。5.4 推理時(shí)預(yù)測(cè)結(jié)果全是同一類現(xiàn)象不管輸入什么圖片模型都輸出“其他垃圾”。原因通常是加載權(quán)重時(shí)沒有正確匹配或者推理預(yù)處理和訓(xùn)練時(shí)不一致。解決先確認(rèn)load_state_dict沒有報(bào)錯(cuò)如果有unexpected keys說明模型結(jié)構(gòu)對(duì)不上然后檢查推理的Normalize參數(shù)是不是和訓(xùn)練時(shí)一樣最后用訓(xùn)練集里的一張圖做推理如果訓(xùn)練集圖片都預(yù)測(cè)錯(cuò)那基本是權(quán)重加載的問題。我一般會(huì)在推理腳本開頭加一句print(model.fc)確認(rèn)分類頭結(jié)構(gòu)。5.5 數(shù)據(jù)加載速度慢導(dǎo)致 GPU 利用率低現(xiàn)象nvidia-smi顯示 GPU 利用率只有 20% 到 30%訓(xùn)練一個(gè) epoch 要很久。原因通常是num_workers設(shè)太小或者圖片尺寸太大。解決把num_workers設(shè)成 CPU 核心數(shù)的一半左右比如 8 核 CPU 設(shè) 4如果圖片原始分辨率是 4000x3000先在預(yù)處理階段統(tǒng)一縮放到 256x256 再存一份能大幅減少 IO 時(shí)間。另外把數(shù)據(jù)放在 SSD 上而不是機(jī)械硬盤這個(gè)提升立竿見影。6. 把項(xiàng)目改出花從及格到優(yōu)秀的三個(gè)進(jìn)階技巧第一個(gè)技巧是換 backbone 做對(duì)比實(shí)驗(yàn)。畢業(yè)設(shè)計(jì)如果只用一個(gè) ResNet18實(shí)驗(yàn)部分會(huì)顯得單薄。你可以把build_model函數(shù)擴(kuò)展一下支持efficientnet_b0、resnet50、mobilenet_v2三種然后在同一份數(shù)據(jù)上跑對(duì)比用表格呈現(xiàn)準(zhǔn)確率、參數(shù)量、推理時(shí)間三個(gè)指標(biāo)。這樣論文的第四章就有了扎實(shí)的內(nèi)容。具體做法是在torchvision.models里找對(duì)應(yīng)的類替換分類頭的代碼邏輯是一樣的注意efficientnet的分類頭叫classifierresnet叫fc改的時(shí)候別搞混。第二個(gè)技巧是加一個(gè)簡(jiǎn)單的 Web 演示界面。用 Gradio 或者 Streamlit 十幾行代碼就能搭起來答辯的時(shí)候現(xiàn)場(chǎng)上傳圖片、實(shí)時(shí)顯示分類結(jié)果和置信度比放 PPT 截圖直觀得多。Gradio 的代碼大概長(zhǎng)這樣import gradio as gr def classify_image(image): label, conf predict(image) return f{label} (置信度: {conf:.2%}) interface gr.Interface( fnclassify_image, inputsgr.Image(typefilepath), outputstext, title垃圾分類識(shí)別系統(tǒng) ) interface.launch(shareFalse)shareFalse表示只在本地局域網(wǎng)訪問答辯時(shí)用自己電腦跑就行。這個(gè)界面不需要任何前端知識(shí)gr.Image自動(dòng)處理上傳和格式轉(zhuǎn)換fn接收文件路徑返回字符串。第三個(gè)技巧是導(dǎo)出 ONNX 模型做推理加速。PyTorch 模型在 CPU 上推理有時(shí)候比較慢轉(zhuǎn)成 ONNX 之后可以用 ONNX Runtime 加速在同樣的硬件上通常能快 1.5 到 2 倍。導(dǎo)出代碼import torch.onnx dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, garbage_cls.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version11 )dynamic_axes讓導(dǎo)出的模型支持變長(zhǎng) batchopset_version11兼容性最好。導(dǎo)出之后用onnxruntime加載推理速度提升明顯而且部署到?jīng)]有 PyTorch 環(huán)境的機(jī)器上也能跑。這三個(gè)技巧我每次做圖像分類項(xiàng)目都會(huì)走一遍尤其是對(duì)比實(shí)驗(yàn)和可視化屬于投入產(chǎn)出比最高的部分。從那以后我每次拿到一個(gè)新的分類項(xiàng)目都強(qiáng)制自己先把 baseline 跑通、再逐步加改進(jìn)絕不一上來就堆模塊。希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取