简介面向深度学习与计算机视觉初学者的垃圾分类实战项目利用卷积神经网络实现生活垃圾图像识别与自动分类可直接用于毕业设计、期末大作业或课程设计也可作为入门计算机视觉的典型练习案例。压缩包约4.21MB体量较小便于快速获取内含完整项目源码及说明文档覆盖数据读取、模型设计、训练评估、预测推理等关键环节下载后配置Python运行环境即可直接使用。目前已有654人浏览/学习关注度持续上升尤其适合希望以完整项目获取高分或快速积累CV实战经验的学生。项目代码结构清晰模块解耦合理不仅便于理解图像分类的标准流程还可支持后续扩展与二次开发帮助掌握数据增强、超参数调节、模型保存与加载等实用技能为深入研究计算机视觉打下扎实基础。1. 垃圾分类识别不是“多分类”这么简单把一张矿泉水瓶照片扔进分类模型正面拍能识别成可回收垃圾换个角度、换个光照就分错这是计算机视觉大作业和毕业设计里最常见的翻车现场。问题根源不在模型不够深而在数据标注方式、类别定义和训练策略没有对齐真实场景。基于深度学习的垃圾分类项目本质不是“做个分类器”而是要把图像采集、标注质量、模型选型和部署推理串成一条完整链路。这篇文章把一个可运行的垃圾分类视觉方案拆开讲覆盖数据集组织、迁移学习选型、训练调参、模型导出和摄像头实时推理每一步都有可直接拷贝的代码适合做期末大作业、课程设计或毕业设计救急也适合想快速上手深度学习图像分类实战的开发者和学生。2. 数据集构建与预处理标注顺序决定训练效率2.1 先把类别定清楚垃圾分类的数据集公开渠道不少常见的有华为云的垃圾分类数据集、Kaggle 上的 Garbage Classification 等类别数从 40 多类到 80 多类不等。但课程设计和毕业设计不建议直接全量训练因为细粒度类别如“塑料瓶”“易拉罐”“玻璃瓶”等彼此视觉相似容易把模型逼到过拟合。常见做法是先把细类聚合为四个大类可回收垃圾、有害垃圾、厨余垃圾、其他垃圾对应现实中的四分类垃圾桶。这样模型容量需求低训练时间快答辩时也更容易解释清楚类别映射逻辑。数据集目录结构建议直接采用ImageFolder可识别的形式减少自定义 Dataset 的调试成本dataset/ ├── train/ │ ├── recycle/ # 可回收 │ ├── hazardous/ # 有害 │ ├── kitchen/ # 厨余 │ └── other/ # 其他 ├── val/ │ ├── recycle/ │ ├── hazardous/ │ ├── kitchen/ │ └── other/ └── test/ ├── recycle/ ├── hazardous/ ├── kitchen/ └── other/每个子文件夹内部不要套多层目录torchvision.datasets.ImageFolder会把顶层文件夹名当作类别标签按字母序排序生成 label。如果四分类目录名分别为hazardous、kitchen、other、recycle则 label 0 对应hazardouslabel 1 对应kitchenlabel 2 对应otherlabel 3 对应recycle。这个隐式顺序容易在后续推理输出中文标签时搞错建议事先用dataset.class_to_idx打印确认。2.2 数据增强别让模型背下背景垃圾分类图像的采集环境差异很大宿舍里拍一张和实验室里拍一张背景完全不同。如果训练集只在单一背景上采集模型很容易把背景纹理当作分类依据实际部署时就崩。数据增强要模拟真实场景下的光照变化、拍摄角度偏移和遮挡常用的组合如下import torchvision.transforms as transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale(0.6, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.3, contrast0.3, 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((224, 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)控制裁剪区域占原图比例这里设置 0.6 到 1.0 是为了避免过度裁剪导致物体主体丢失。RandomRotation(15)只旋转 15 度超过这个角度会出现大面积黑边反而引入噪声。ColorJitter的三个参数分别扰动亮度、对比度和饱和度模拟不同光照环境。验证集和测试集不参与任何随机增强只做 Resize 和归一化这样才能给模型一个稳定的评估基线。归一化用的 mean 和 std 是 ImageNet 统计值因为后面迁移学习加载的是 ImageNet 预训练权重输入分布必须与预训练时一致。2.3 用 DataLoader 组织数据流from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset ImageFolder(rootdataset/train, transformtrain_transform) val_dataset ImageFolder(rootdataset/val, transformval_transform) train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers2, pin_memoryTrue )batch_size选 32 是兼顾显存和 BatchNorm 统计稳定性的默认值如果 GPU 显存只有 4GB 可以降到 16。num_workers在 Windows 上设置为 0 或 2Linux 可以开到 4 以上Windows 下 worker 数量过高容易触发 DataLoader 报错。pin_memoryTrue适合 GPU 训练能减少主机到设备的数据拷贝时间。训练集shuffleTrue打乱样本顺序避免模型按批次顺序学到无关规律验证集shuffleFalse保证每次评估顺序一致便于对比不同 epoch 的指标差异。3. 模型选型与迁移学习ResNet50 和 MobileNetV3 怎么挑3.1 为什么要用迁移学习垃圾分类公开数据集的规模通常在几千到几万张而 ImageNet 有 120 万张以上。从零训练一个 CNN在小数据集上极易陷入过拟合具体表现是训练准确率接近 100%验证准确率却只有六七成。迁移学习的思路是加载在 ImageNet 上预训练好的权重这些权重已经学会了边缘、纹理、颜色组合等通用视觉特征我们只需要在特定数据集上微调后几层即可。这在垃圾分类场景下尤其有效因为瓶子的轮廓、纸张的纹理等底层特征和自然图像是通用的而“可回收”“厨余”这种高层语义差异正好适合微调阶段去学习。import torchvision.models as models import torch.nn as nn model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) for param in model.parameters(): param.requires_grad False num_features model.fc.in_features model.fc nn.Linear(num_features, 4) for param in model.fc.parameters(): param.requires_grad True这里requires_grad False表示冻结 backbone 参数只训练新替换的分类头。model.fc.in_features是 ResNet50 最后一个池化层输出的特征维度值为 2048。分类头输出维度从 1000 改成 4对应四个垃圾类别。先冻结训练的是策略是分类头随机初始化梯度噪声大如果 backbone 同时更新容易破坏 ImageNet 预训练特征先训练分类头几轮再解锁 backbone 低学习率微调实践中收敛更稳。3.2 两个模型的参数对比和选型依据模型参数量输入尺寸单张推理时延CPU适用场景ResNet5025.6M224x224约 40msGPU 训练、准确率优先MobileNetV3-Large5.4M224x224约 10ms边缘部署、CPU 推理ResNet50 的残差结构通过跳跃连接缓解了深层网络的梯度消失问题在 ImageNet 上表现稳定是课程设计和期末大作业的主流选择导师对它的结构熟悉度也高答辩不容易被追问到冷门细节。MobileNetV3 则用深度可分离卷积大幅压缩参数量如果机器没有独立 GPU或者想部署到树莓派、Jetson Nano 等设备上优先选它。需要说明的是小模型的泛化能力不一定弱于大模型尤其是数据集规模不到一万张时MobileNetV3 的归纳偏置反而可能带来更好的验证集表现。我的习惯是先跑 ResNet50 拿到一个可靠的 baseline再切 MobileNetV3 对比精度和速度选择能跑得动且精度达标的那一个。3.3 替换分类头的代码细节import torch.nn as nn import torchvision.models as models def build_model(model_name: str, num_classes: int 4, pretrained: bool True): if model_name resnet50: weights models.ResNet50_Weights.IMAGENET1K_V1 if pretrained else None model models.resnet50(weightsweights) model.fc nn.Sequential( nn.Dropout(p0.2), nn.Linear(model.fc.in_features, 256), nn.ReLU(inplaceTrue), nn.Linear(256, num_classes) ) elif model_name mobilenetv3: weights models.MobileNetV3_Large_Weights.IMAGENET1K_V1 if pretrained else None model models.mobilenet_v3_large(weightsweights) model.classifier[-1] nn.Linear(model.classifier[-1].in_features, num_classes) else: raise ValueError(fUnsupported model: {model_name}) return model分类头可以不是单层 Linear加入 Dropout 和中间层能提升分类头容量但也增加了过拟合风险。这里的Dropout(p0.2)在训练时随机丢弃 20% 的神经元推理时自动关闭相当于对分类头做了集成。MobileNetV3 的classifier是 Sequential 结构最后一层是 Linear直接替换classifier[-1]保留前面已有的 Dropout 层。两种模型在冻结 backbone 时都需要确保只对分类头的参数计算梯度否则优化器会把所有参数的梯度都算一遍冻结设置就白做了。4. 训练管线与调参从过拟合到收敛4.1 损失函数与优化器配置四分类问题用交叉熵损失即可nn.CrossEntropyLoss()内部已经包含 Softmax不需要在网络输出层额外加 Softmax。优化器方面微调阶段不要用 SGD 默认参数直接跑ResNet 系列在数据量小时容易震荡我一般用 AdamW权重衰减设到1e-4到1e-3之间作为正则化手段。import torch import torch.nn as nn criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW( [p for p in model.parameters() if p.requires_grad], lr3e-4, weight_decay1e-4 )AdamW 的weight_decay实现了解耦权重衰减相比 Adam 的 L2 正则化在 Transformer 和 CNN 上都有更稳定的收敛表现。学习率初始值3e-4适合冻结 backbone 后只训练分类头的阶段如果整个网络解冻需要降到1e-4以下否则 backbone 的预训练特征容易被大步长更新破坏。4.2 余弦退火比固定学习率收敛得更稳固定学习率训练在 epoch 后期容易在损失曲面底部来回震荡余弦退火让学习率随训练进度从初始值平滑降低到接近 0前中期保持较大步长快速收敛后期以微小步长在最优解附近精细调整。配合T_max设置为总 epoch 数效果立竿见影。total_epochs 30 scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxtotal_epochs, eta_min1e-6 )T_max30表示在 30 个 epoch 内完成一个完整余弦周期eta_min是最低学习率。如果训练只跑 15 个 epoch 就停了学习率还没降到足够低模型可能仍然处于损失较大区域所以设置T_max要和实际训练的 epoch 数一致或者设置一个稍大的值配合 Early Stopping。4.3 类别不平衡用 WeightedRandomSampler 治标垃圾分类四个类别在公开数据集中数量并不均等常见的是“其他垃圾”图片最多而“有害垃圾”因为收集门槛高图片数量明显偏少。如果直接用原始分布训练模型对少数类的召回率会很差。处理方式有两种一是使用WeightedRandomSampler在采样时让样本量少的类别有更高概率被抽到二是调整损失权重给少数类的交叉熵损失乘以更大系数。import torch from torch.utils.data import WeightedRandomSampler labels [train_dataset.targets[i] for i in range(len(train_dataset))] class_counts torch.bincount(torch.tensor(labels)).float() class_weights 1.0 / class_counts sample_weights class_weights[labels] train_sampler WeightedRandomSampler( weightssample_weights, num_sampleslen(sample_weights), replacementTrue ) train_loader DataLoader( train_dataset, batch_size32, samplertrain_sampler, num_workers4, pin_memoryTrue )每个样本的权重是1 / class_count类别样本越少权重越大。replacementTrue表示抽样放回保证每个 batch 里都可能出现少数类样本。使用自定义 sampler 时DataLoader 的shuffle参数必须设为False因为 sampler 已经承担了打乱顺序的职责。这个方案能改善少数类的训练覆盖但不会凭空增加信息量效果上限还是取决于原始数据本身。4.4 训练主循环与验证指标跟踪def train_one_epoch(model, loader, criterion, optimizer, device): model.train() running_loss, correct, total 0.0, 0, 0 for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) avg_loss running_loss / total accuracy correct / total return avg_loss, accuracy def validate(model, loader, criterion, device): model.eval() running_loss, correct, total 0.0, 0, 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) avg_loss running_loss / total accuracy correct / total return avg_loss, accuracy训练和验证逻辑的关键差异在model.train()和model.eval()的切换。train()模式启用 Dropouteval()模式关闭 Dropout 并且 BatchNorm 使用运行均值而不是当前 batch 统计量。如果验证时忘记切换eval()BatchNorm 的统计值会被验证集污染导致训练过程出现莫名其妙的精度波动。torch.max(outputs, 1)返回每个样本的最大概率值和对应索引索引就是预测类别。建议每个 epoch 结束后打印训练损失、训练准确率、验证损失、验证准确率观察训练损失下降但验证损失上升的节点那就是过拟合开始的位置。5. 推理部署TorchScript 导出与实时识别5.1 把模型固化为 TorchScriptPyTorch 的torch.jit.script可以把训练好的nn.Module编译为独立可序列化的 TorchScript 模型部署时不需要依赖 Python 环境和 PyTorch 完整框架直接用torch.jit.load加载即可。相比直接保存state_dictTorchScript 模型内嵌了网络结构避免了推理端重新定义模型类的麻烦import torch model.eval() example_input torch.randn(1, 3, 224, 224) scripted_model torch.jit.trace(model, example_input) scripted_model.save(garbage_classifier.pt) print(Model exported to garbage_classifier.pt)torch.jit.trace需要传入一个固定的示例输入沿着输入数据走一遍模型计算图并记录所有操作所以输入尺寸必须是固定的(1, 3, 224, 224)通道格式。如果模型里包含数据相关的控制流比如根据 batch 大小走不同分支trace 可能无法完整覆盖垃圾分类模型是纯卷积加全连接结构不存在这个问题。批量大小固定为 1 的部署场景使用 trace 最稳妥。5.2 单张图片推理流程import torch from PIL import Image import torchvision.transforms as transforms device torch.device(cuda if torch.cuda.is_available() else cpu) model torch.jit.load(garbage_classifier.pt, map_locationdevice) model.to(device).eval() class_names [hazardous, kitchen, other, recycle] class_names_cn [有害垃圾, 厨余垃圾, 其他垃圾, 可回收垃圾] inference_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def predict_image(image_path: str, top_k: int 2): image Image.open(image_path).convert(RGB) tensor inference_transform(image).unsqueeze(0).to(device) with torch.no_grad(): logits model(tensor) probabilities torch.softmax(logits, dim1) top_prob, top_idx torch.topk(probabilities, top_k, dim1) top_prob top_prob.squeeze(0).tolist() top_idx top_idx.squeeze(0).tolist() for score, idx in zip(top_prob, top_idx): print(f{class_names_cn[idx]:6s} ({class_names[idx]}): {score*100:.1f}%) predict_image(test/kitchen_001.jpg)image.convert(RGB)很重要因为有些 JPEG 图片是灰度模式或带透明通道的 RGBA 格式直接送入模型会因通道数不匹配报错。unsqueeze(0)把(3, 224, 224)扩展成(1, 3, 224, 224)的 batch 维度。torch.topk返回前 k 个最大概率值和索引展示前两个结果可以让用户知道模型是否有把握如果前两个概率非常接近说明图像存在歧义。这里的 class_names 顺序必须和训练时dataset.class_to_idx的排序一致否则输出标签会错位。5.3 摄像头实时垃圾分类识别import cv2 import torch import torchvision.transforms as transforms from PIL import Image device torch.device(cuda if torch.cuda.is_available() else cpu) model torch.jit.load(garbage_classifier.pt, map_locationdevice) model.to(device).eval() class_names_cn [有害垃圾, 厨余垃圾, 其他垃圾, 可回收垃圾] cap cv2.VideoCapture(0) if not cap.isOpened(): print(无法打开摄像头) exit(1) while True: ret, frame cap.read() if not ret: break pil_image Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) input_tensor inference_transform(pil_image).unsqueeze(0).to(device) with torch.no_grad(): output torch.softmax(model(input_tensor), dim1) score, idx torch.max(output, dim1) label class_names_cn[idx.item()] cv2.putText(frame, f{label} {score.item()*100:.1f}%, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 2) cv2.imshow(Garbage Classification, frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows()OpenCV 读取的帧是 BGR 顺序PIL 的Image.fromarray默认按 RGB 处理所以必须先转换颜色空间否则画面会出现明显的蓝色偏色推理准确率也大幅下降。摄像头推理的瓶颈在预处理和推理耗时ResNet50 在普通 CPU 上大约 40 毫秒一帧加上 OpenCV 显示延迟基本能到 20 FPS 左右如果帧率不理想可以先把Resize目标降到(160, 160)观察精度损失。cv2.waitKey(1)的返回值是用户按键的 ASCII 码这里用按q退出循环。注意如果摄像头不是默认设备VideoCapture(0)的参数要改为对应设备编号。6. 答辩加分用 Grad-CAM 可视化模型关注区域6.1 原理与在垃圾分类中的意义Grad-CAM 通过计算类别得分对最后一个卷积层特征图的梯度得到每个特征通道的重要性权重再对特征图加权求和生成热力图。热力图叠加到原图上就能直观看到模型把图像中的哪些区域作为分类依据。例如模型把“易拉罐”识别成“可回收垃圾”热力图应当集中在易拉罐本身而不是背景桌面如果热力图集中在背景区域说明模型学到了环境特征需要回到数据增强部分做调整。这一招在期末大作业答辩中非常出彩因为大部分同学的演示只停留在准确率数字上能讲清楚模型关注区域体现了对深度学习方法内在机制的理解。6.2 用 Hook 捕获中间层特征图import torch import torch.nn.functional as F import numpy as np import cv2 from PIL import Image def grad_cam(model, input_tensor, target_layer, class_idx): feature_maps {} gradients {} def forward_hook(module, input, output): feature_maps[value] output def backward_hook(module, grad_input, grad_output): gradients[value] grad_output[0] hook1 target_layer.register_forward_hook(forward_hook) hook2 target_layer.register_full_backward_hook(backward_hook) model.eval() output model(input_tensor) score output[0, class_idx] model.zero_grad() score.backward() activations feature_maps[value].squeeze(0) grads gradients[value].squeeze(0) weights torch.mean(grads, dim(1, 2)) cam torch.zeros(activations.shape[1:], dtypetorch.float32) for i in range(activations.shape[0]): cam weights[i] * activations[i] cam F.relu(cam) cam cam - cam.min() cam cam / (cam.max() 1e-8) hook1.remove() hook2.remove() return cam.detach().cpu().numpy()register_forward_hook在前向传播结束时捕获该层输出register_full_backward_hook捕获反向传播时该层接收到的梯度。register_full_backward_hook接收的参数有三个grad_output[0]才是目标层输出的梯度形状。score output[0, class_idx]选定某个类别对其 backward 得到针对该类别的梯度。热力图可视化的关键是torch.mean(grads, dim(1, 2))它对每个通道的梯度做全局平均池化得到通道权重。F.relu(cam)只保留正贡献区域丢弃对分类起到抑制作用的特征。归一化时加上1e-8避免除零。6.3 热力图叠加与答辩展示建议def overlay_heatmap(image_path, cam, alpha0.5): image cv2.imread(image_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image cv2.resize(image, (224, 224)) heatmap cv2.resize(cam, (image.shape[1], image.shape[0])) heatmap np.uint8(255 * heatmap) heatmap cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) overlay cv2.addWeighted(image, 1 - alpha, heatmap, alpha, 0) return overlayalpha0.5一般视觉效果最好热力图颜色太浓会盖住物体轮廓太淡又看不出关注区域。展示时把原图、热力图、叠加图三张图横向拼在一起配上预测概率值一起输出。答辩中遇到导师追问”模型为什么判断这张是可回收垃圾“直接给出叠加图并指出热力集中区域说服力远大于纯口头解释。值得注意的是Grad-CAM 的score.backward()每次只针对一个类别如果要展示所有类别就需要对每个类别分别计算多类别多图像的批量可视化在课堂上演示时建议提前离线生成好图片避免现场等待算力热身。此外ResNet50 最后一层卷积层的输出分辨率是 7x7热力图放大到原图尺寸后边缘会比较粗糙这是正常现象不需要改用更高分辨率层。本文还有配套的精品资源点击获取