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