简介这是一份基于深度学习的垃圾分类项目工程包适用于高校期末大作业、本科毕业设计或机器学习入门实践。项目以Python为主要开发语言覆盖图像预处理、数据集处理、模型构建、训练评估与简单应用部署的完整流程适合希望快速搭建图像分类系统并理解工程结构的学习者。压缩包共43个文件以17个py源码文件为核心包括模型定义、数据预处理、数据加载及启动入口等模块另含15个pyc编译文件、5个xml工程配置、3个ipynb交互式笔记和readme说明文档整体仅55KB轻量易用。models目录内置squeezenet、alexnet、inception、densenet、resnet、vgg等多种经典网络实现便于对比与替换utils目录提供日志、评估、JSON处理等工具函数工程结构清晰。目前已有45人学习下载适合深度学习初学者借鉴项目组织方式也可作为垃圾分类课题的基础代码框架进行二次开发。1. 垃圾分类这个场景为什么值得把深度学习从头落地一遍垃圾分类和图像分类不太一样类别之间边界模糊塑料瓶和玻璃瓶远看都是“透明圆柱体”厨余垃圾里的菜叶和树叶在颜色纹理上几乎无法用人工规则区分。小区试点里常见的问题是摄像头装了、模型也跑了但识别准确率停留在 70% 上下一到晚上灯光偏黄时就开始乱报。深度学习方法在这里的优势不是“用更深的网络刷精度”而是通过数据增强和迁移学习把模型对光照、角度、遮挡的鲁棒性提前做进去。从工程角度看一个打包成 zip 的“基于深度学习的垃圾分类”项目通常包含四类东西数据集、训练代码、模型权重、部署脚本。这套组合能解决的问题也明确如何在有限算力下让模型在几百类垃圾上做到可用精度如何把训练好的模型导出成轻量格式放到摄像头或边缘盒子上。适合的人群是刚做完图像分类教程、想迁移到真实场景的工程师以及准备做毕设或比赛但不想只跑通一个 ResNet 的学生。2. 拿到 zip 包先做的事数据集解构、标签整理与 DataLoader2.1 解压后先看清目录结构train/val/test 的划分方式决定后面所有事常见的垃圾分类数据集打包结构有两种。一种是按类别分文件夹形如train/plastic/001.jpgtrain/cardboard/002.jpg另一种是 CSV 标注文件加平铺图片目录。前者对应torchvision.datasets.ImageFolder后者需要自己写 Dataset。第一步永远是打开 zip 后先tree一眼不要急着写训练脚本。unzip garbage_dataset.zip -d ./data find ./data -maxdepth 2 -type d | head -30 du -sh ./data/*.zip 2/dev/null || trueunzip解压后先确认目录深度因为很多数据集打包时多套了一层外层目录ImageFolder会把这个外层目录当成一个类别训练时准确率直接崩掉。du -sh是看有没有重复打包我遇到过两次数据集里同时存在原始图片和增强后图片导致验证集和训练集内容重叠指标虚高到 99%换到真实场景立刻现原形。2.2 用 ImageFolder 还是自定义 Dataset一个带验证集划分的装载代码如果目录结构规整ImageFolder就够了。但垃圾分类数据集的标注质量参差不齐常见问题包括同一种垃圾出现在两个类别文件夹里、部分图片是损坏的零字节文件、少数类别图片数量只有个位数。这些问题用ImageFolder看不出来得在装载阶段把类别数量、单类图片数打印出来。from torchvision import datasets, transforms from torch.utils.data import DataLoader, random_split import torch train_dir ./data/train full_ds datasets.ImageFolder(train_dir) print(f类别总数: {len(full_ds.classes)}) print(f图片总数: {len(full_ds)}) for cls, idx in full_ds.class_to_idx.items(): count sum(1 for _, label in full_ds.samples if label idx) print(f{cls}: {count} 张) train_size int(len(full_ds) * 0.85) val_size len(full_ds) - train_size train_ds, val_ds random_split(full_ds, [train_size, val_size]) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4)这段代码里的class_to_idx是 zip 包项目里最容易被忽略的资产。垃圾分类数据集的类别顺序通常和训练时的文件夹排列顺序一致但不同数据集的排序不一致。训练结束后做推理时要从train_ds.dataset.classes里导出类别映射表保存成 JSON 或 Python 文件后续部署脚本直接读这个映射避免“模型输出索引 5但不知道索引 5 是什么垃圾”的尴尬。2.3 数据增强参数表训练集和验证集必须分开配置垃圾分类识别里最有效的增强不是随机的旋转而是光照扰动和遮挡模拟。因为实际场景里摄像头装在回收箱上方视角基本固定但阳光角度、灯光色温和人伸手遮挡的情况非常多。训练集和验证集的增强策略按下面这套配置增强操作训练集验证集参数说明RandomResizedCrop是仅Resize裁剪范围 0.61.0模拟远近变化ColorJitter是否brightness0.4, contrast0.3, saturation0.2RandomHorizontalFlip是否概率 0.5垃圾方向不影响类别RandomErasing是否p0.25模拟遮挡Normalize是是mean[0.485,0.456,0.406], std[0.229,0.224,0.225]验证集不随机化这是很多刚上手项目的人常犯的错。验证集如果也做随机裁剪同一次验证跑出来的准确率波动很大你无法判断是模型变好了还是裁剪碰巧取到了关键区域。垃圾分类数据里有一类很典型的情况矿泉水瓶和易拉罐在裁剪区域只剩瓶身局部时颜色是唯一的判别线索。所以验证集统一先 Resize 到 256再 CenterCrop 到 224保证每次评测的输入一致性。3. 从迁移学习到模型收敛分类模型选型与训练参数3.1 垃圾分类场景下四个候选模型的精度与参数量对比垃圾分类是典型的细粒度图像分类类间相似度高、类内差异也不小。模型选型上不需要追求 SOTA但要让算力有限的边缘设备跑得动。四个常见候选模型的对比参数如表所示模型Top-1 精度ImageNet参数量单张推理耗时CPU适用场景ResNet1869.8%11.7M约 18ms快速验证基线ResNet5076.1%25.6M约 42ms准确率优先有 GPUMobileNetV271.8%3.5M约 8ms边缘设备首选EfficientNet-B077.1%5.3M约 14ms均衡型选择垃圾分类数据集里通常只有 40 到 80 个类别比 ImageNet 的 1000 类简单得多。因此 ResNet18 已经能跑出相当好的基线MobileNetV2 在精度掉 1 到 2 个点的前提下推理速度能提升一倍以上。如果 zip 包里带显卡驱动环境配置说明一般默认是 NVIDIA GPU CUDA没有的话先用 MobileNetV2 在 CPU 上做完整流程验证。3.2 在 ImageNet 权重上做微调冻结前半部分的实操代码垃圾分类数据集再大也就几万张图片直接从零训练容易过拟合。标准做法是用 torchvision 里预训练好的权重做迁移学习。关键选择是冻结多少层冻结太多模型学不到垃圾的特殊纹理冻结太少小数据集上微调容易把预训练特征冲掉。常用的折中是冻结前 60% 的层只训练后半段和分类头。import torchvision.models as models import torch.nn as nn model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) # 冻结卷积层的前半部分 params_to_freeze list(model.parameters())[:60] for param in params_to_freeze: param.requires_grad False optimizer torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr3e-4, weight_decay1e-4 ) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience3 )list(model.parameters())[:60]这个写法依赖模型内部参数顺序ResNet18 里前 60 个参数基本覆盖到 layer2 之前。换模型时需要先打印参数名再调整。Adam 的初始学习率 3e-4 是迁移学习里经验最稳的起点不要上来就用 1e-3微调阶段大学习率会直接破坏预训练权重。weight_decay 设 1e-4 是为了抑制过拟合垃圾分类类别数少模型容量容易过剩。3.3 训练循环里的实际参数学习率、batch_size 与类别不均衡处理import torch.nn.functional as F best_acc 0.0 for epoch in range(30): model.train() train_loss 0.0 for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss F.cross_entropy(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() model.eval() correct, total 0, 0 with torch.no_grad(): for images, labels in val_loader: outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total scheduler.step(val_loss) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), ./best_model.pth) print(fEpoch {epoch1}: loss{train_loss:.4f}, val_acc{val_acc:.4f})这里没有直接用torch.optim.lr_scheduler.CosineAnnealingLR而是用ReduceLROnPlateau因为垃圾分类数据集的收敛曲线不像标准分类任务那么平滑初始学习率偏大时验证损失会上下震荡plateau 模式会在连续 3 个 epoch 不下降时自动减半。训练时把best_model.pth和最后一次 epoch 的权重分开保存因为轻微过拟合状态下验证准确率最高这个权重用于后续推理往往比最终权重效果好。如果数据集里“纸箱”有一万张而“电池”只有两百张交叉熵损失会偏向多数类。常见的做法是把F.cross_entropy的weight参数设为各类样本数的倒数class_counts torch.tensor([...]) class_weights 1.0 / class_counts.float() class_weights class_weights / class_weights.sum() * num_classes criterion nn.CrossEntropyLoss(weightclass_weights)加了类别权重后多数类的梯度会被压低少数类的错误预测会被放大惩罚。代价是训练集 loss 偏大、验证准确率可能微降但混淆矩阵里少数类的召回率会明显提升。这个取舍在垃圾分类场景里通常是值得的因为有害垃圾和电池这类少数类才是分错成本最高的。4. 从混淆矩阵到 ONNX 导出识别性能的验证与部署形态4.1 只看准确率会漏掉的错误用混淆矩阵定位“类间混淆对”训练结束不要只记录 val_acc要单独跑一次完整验证集并画出混淆矩阵。垃圾分类项目里准确率 90% 和 92% 的模型差异往往不来自某一个类别的整体提升而是几对“易混淆类”的处理方式变了。打印出错误最多的前 10 对组合from sklearn.metrics import confusion_matrix import numpy as np all_preds, all_labels [], [] model.eval() with torch.no_grad(): for images, labels in val_loader: outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) np.fill_diagonal(cm, 0) sorted_pairs np.dstack(np.unravel_index(np.argsort(cm.ravel())[::-1], cm.shape))[0][:10] for pred_cls, true_cls in sorted_pairs: if cm[true_cls][pred_cls] 0: print(f真实类别 {classes[true_cls]} 被误判为 {classes[pred_cls]}: {cm[true_cls][pred_cls]} 次)np.fill_diagonal(cm, 0)把对角线清零后剩下的非零位置全部是错误配对。输出这些对之后去看对应图片的样本通常会发现两个规律要么是这两类垃圾在颜色或形状上高度相似比如干净的酸奶盒和白纸要么是数据集中这两类的拍摄背景重叠太多模型实际学的是背景特征而不是垃圾本身。后者的解决思路是把背景变化类似的两类样本重新整理或者增加针对性增强。4.2 模型压缩与边缘部署的一条快捷路径ONNX 导出 量化训练好的 PyTorch 权重不能直接放到边缘盒子上跑常见部署链路是先转 ONNX再根据目标设备选择精度格式。import torch.onnx model.load_state_dict(torch.load(./best_model.pth)) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, ./garbage_model.onnx, input_names[input], output_names[logits], dynamic_axes{input: {0: batch_size}}, opset_version11 )导出后检查 ONNX 文件的输入输出维度确认最后一层输出是类别数而不是 1000。很多压缩包项目会在导出这一步出错原因是用torchvision.models.resnet18()加载了未修改分类头的预训练模型模型没有换成num_classes的分类层。导出改成动态 batch 尺寸后同一个模型可以适配摄像头单帧输入和批量评测。量化方面分类模型的常用策略压缩方式精度影响推理加速部署方式动态量化INT8分类任务约 0.5% 损失CPU 提升 1.52 倍torch.quantization.quantize_dynamic静态量化INT8需校准集损失可控制在 1% 内CPU 提升 23 倍需要少量真实场景图片做校准ONNX Runtime FP16几乎无损失GPU 提升明显适合 Jetson 系列垃圾分类部署最常见的两个场景是树莓派/工控机上的 CPU 推理和 Jetson Nano 上的轻量 GPU 推理。前者用动态量化就够后者导出 FP16 ONNX 跑 TensorRT 更合适。5. 垃圾分类模型训练完最值得记下的三个工程细节5.1 类别映射表 JSON 化部署时最容易被忽视的坑训练用的class_to_idx只在训练代码里存在部署时常常被丢掉。离线推理时如果手里只有 ONNX 模型和 numpy 数组没有类别名输出就是一个数字。把类别映射在训练最后一步单独保存和权重文件放一起import json with open(./class_map.json, w) as f: json.dump(full_ds.classes, f, ensure_asciiFalse)这里直接用列表顺序对应索引即可json.load读出来也是有序的保证推理代码里pred_cls class_names[pred_idx]不会错位。这个 JSON 文件和best_model.pth、ONNX 文件建议打成同一个压缩包避免分发时漏文件。5.2 早停策略的两个参数patience 和 min_delta垃圾分类数据集的验证准确率通常在 epoch 10 到 15 之间进入平台期继续训练只会过拟合。早停的参数设置经验是 patience 取 5 到 7min_delta 取 0.001。min_delta 的作用是忽略精度波动比如验证 loss 从 1.23 到 1.22 这种变化不算“变好”避免模型保存过于频繁。配合ReduceLROnPlateau一起用时先降学习率再早停比直接停的效果好。best_val_loss float(inf) epochs_no_improve 0 patience 5 for epoch in range(30): val_loss validate(model, val_loader) if val_loss best_val_loss - 0.001: best_val_loss val_loss epochs_no_improve 0 torch.save(model.state_dict(), ./best_model.pth) else: epochs_no_improve 1 if epochs_no_improve patience: print(fEarly stopping at epoch {epoch}) break5.3 测试时增强验证集上多跑几次取平均预测TTA 是把同一张图片做多次不同预处理后取平均输出。对单张图片分别做 CenterCrop 和四角裁剪五张图的结果做 softmax 平均可以在不训练的情况下把准确率提升 0.5 到 1 个百分点。实现时注意 batch 维度堆叠把五次增强结果一起 forward然后torch.mean(torch.softmax(logits, dim1), dim0)取平均。测试时增强会多花五倍推理时间适合对单张照片精度有要求的场景实时视频流里基本不用。本文还有配套的精品资源点击获取 从迁移学习到ONNX部署：垃圾分类深度学习落地指南
简介这是一份基于深度学习的垃圾分类项目工程包适用于高校期末大作业、本科毕业设计或机器学习入门实践。项目以Python为主要开发语言覆盖图像预处理、数据集处理、模型构建、训练评估与简单应用部署的完整流程适合希望快速搭建图像分类系统并理解工程结构的学习者。压缩包共43个文件以17个py源码文件为核心包括模型定义、数据预处理、数据加载及启动入口等模块另含15个pyc编译文件、5个xml工程配置、3个ipynb交互式笔记和readme说明文档整体仅55KB轻量易用。models目录内置squeezenet、alexnet、inception、densenet、resnet、vgg等多种经典网络实现便于对比与替换utils目录提供日志、评估、JSON处理等工具函数工程结构清晰。目前已有45人学习下载适合深度学习初学者借鉴项目组织方式也可作为垃圾分类课题的基础代码框架进行二次开发。1. 垃圾分类这个场景为什么值得把深度学习从头落地一遍垃圾分类和图像分类不太一样类别之间边界模糊塑料瓶和玻璃瓶远看都是“透明圆柱体”厨余垃圾里的菜叶和树叶在颜色纹理上几乎无法用人工规则区分。小区试点里常见的问题是摄像头装了、模型也跑了但识别准确率停留在 70% 上下一到晚上灯光偏黄时就开始乱报。深度学习方法在这里的优势不是“用更深的网络刷精度”而是通过数据增强和迁移学习把模型对光照、角度、遮挡的鲁棒性提前做进去。从工程角度看一个打包成 zip 的“基于深度学习的垃圾分类”项目通常包含四类东西数据集、训练代码、模型权重、部署脚本。这套组合能解决的问题也明确如何在有限算力下让模型在几百类垃圾上做到可用精度如何把训练好的模型导出成轻量格式放到摄像头或边缘盒子上。适合的人群是刚做完图像分类教程、想迁移到真实场景的工程师以及准备做毕设或比赛但不想只跑通一个 ResNet 的学生。2. 拿到 zip 包先做的事数据集解构、标签整理与 DataLoader2.1 解压后先看清目录结构train/val/test 的划分方式决定后面所有事常见的垃圾分类数据集打包结构有两种。一种是按类别分文件夹形如train/plastic/001.jpgtrain/cardboard/002.jpg另一种是 CSV 标注文件加平铺图片目录。前者对应torchvision.datasets.ImageFolder后者需要自己写 Dataset。第一步永远是打开 zip 后先tree一眼不要急着写训练脚本。unzip garbage_dataset.zip -d ./data find ./data -maxdepth 2 -type d | head -30 du -sh ./data/*.zip 2/dev/null || trueunzip解压后先确认目录深度因为很多数据集打包时多套了一层外层目录ImageFolder会把这个外层目录当成一个类别训练时准确率直接崩掉。du -sh是看有没有重复打包我遇到过两次数据集里同时存在原始图片和增强后图片导致验证集和训练集内容重叠指标虚高到 99%换到真实场景立刻现原形。2.2 用 ImageFolder 还是自定义 Dataset一个带验证集划分的装载代码如果目录结构规整ImageFolder就够了。但垃圾分类数据集的标注质量参差不齐常见问题包括同一种垃圾出现在两个类别文件夹里、部分图片是损坏的零字节文件、少数类别图片数量只有个位数。这些问题用ImageFolder看不出来得在装载阶段把类别数量、单类图片数打印出来。from torchvision import datasets, transforms from torch.utils.data import DataLoader, random_split import torch train_dir ./data/train full_ds datasets.ImageFolder(train_dir) print(f类别总数: {len(full_ds.classes)}) print(f图片总数: {len(full_ds)}) for cls, idx in full_ds.class_to_idx.items(): count sum(1 for _, label in full_ds.samples if label idx) print(f{cls}: {count} 张) train_size int(len(full_ds) * 0.85) val_size len(full_ds) - train_size train_ds, val_ds random_split(full_ds, [train_size, val_size]) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4)这段代码里的class_to_idx是 zip 包项目里最容易被忽略的资产。垃圾分类数据集的类别顺序通常和训练时的文件夹排列顺序一致但不同数据集的排序不一致。训练结束后做推理时要从train_ds.dataset.classes里导出类别映射表保存成 JSON 或 Python 文件后续部署脚本直接读这个映射避免“模型输出索引 5但不知道索引 5 是什么垃圾”的尴尬。2.3 数据增强参数表训练集和验证集必须分开配置垃圾分类识别里最有效的增强不是随机的旋转而是光照扰动和遮挡模拟。因为实际场景里摄像头装在回收箱上方视角基本固定但阳光角度、灯光色温和人伸手遮挡的情况非常多。训练集和验证集的增强策略按下面这套配置增强操作训练集验证集参数说明RandomResizedCrop是仅Resize裁剪范围 0.61.0模拟远近变化ColorJitter是否brightness0.4, contrast0.3, saturation0.2RandomHorizontalFlip是否概率 0.5垃圾方向不影响类别RandomErasing是否p0.25模拟遮挡Normalize是是mean[0.485,0.456,0.406], std[0.229,0.224,0.225]验证集不随机化这是很多刚上手项目的人常犯的错。验证集如果也做随机裁剪同一次验证跑出来的准确率波动很大你无法判断是模型变好了还是裁剪碰巧取到了关键区域。垃圾分类数据里有一类很典型的情况矿泉水瓶和易拉罐在裁剪区域只剩瓶身局部时颜色是唯一的判别线索。所以验证集统一先 Resize 到 256再 CenterCrop 到 224保证每次评测的输入一致性。3. 从迁移学习到模型收敛分类模型选型与训练参数3.1 垃圾分类场景下四个候选模型的精度与参数量对比垃圾分类是典型的细粒度图像分类类间相似度高、类内差异也不小。模型选型上不需要追求 SOTA但要让算力有限的边缘设备跑得动。四个常见候选模型的对比参数如表所示模型Top-1 精度ImageNet参数量单张推理耗时CPU适用场景ResNet1869.8%11.7M约 18ms快速验证基线ResNet5076.1%25.6M约 42ms准确率优先有 GPUMobileNetV271.8%3.5M约 8ms边缘设备首选EfficientNet-B077.1%5.3M约 14ms均衡型选择垃圾分类数据集里通常只有 40 到 80 个类别比 ImageNet 的 1000 类简单得多。因此 ResNet18 已经能跑出相当好的基线MobileNetV2 在精度掉 1 到 2 个点的前提下推理速度能提升一倍以上。如果 zip 包里带显卡驱动环境配置说明一般默认是 NVIDIA GPU CUDA没有的话先用 MobileNetV2 在 CPU 上做完整流程验证。3.2 在 ImageNet 权重上做微调冻结前半部分的实操代码垃圾分类数据集再大也就几万张图片直接从零训练容易过拟合。标准做法是用 torchvision 里预训练好的权重做迁移学习。关键选择是冻结多少层冻结太多模型学不到垃圾的特殊纹理冻结太少小数据集上微调容易把预训练特征冲掉。常用的折中是冻结前 60% 的层只训练后半段和分类头。import torchvision.models as models import torch.nn as nn model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) # 冻结卷积层的前半部分 params_to_freeze list(model.parameters())[:60] for param in params_to_freeze: param.requires_grad False optimizer torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr3e-4, weight_decay1e-4 ) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience3 )list(model.parameters())[:60]这个写法依赖模型内部参数顺序ResNet18 里前 60 个参数基本覆盖到 layer2 之前。换模型时需要先打印参数名再调整。Adam 的初始学习率 3e-4 是迁移学习里经验最稳的起点不要上来就用 1e-3微调阶段大学习率会直接破坏预训练权重。weight_decay 设 1e-4 是为了抑制过拟合垃圾分类类别数少模型容量容易过剩。3.3 训练循环里的实际参数学习率、batch_size 与类别不均衡处理import torch.nn.functional as F best_acc 0.0 for epoch in range(30): model.train() train_loss 0.0 for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss F.cross_entropy(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() model.eval() correct, total 0, 0 with torch.no_grad(): for images, labels in val_loader: outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total scheduler.step(val_loss) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), ./best_model.pth) print(fEpoch {epoch1}: loss{train_loss:.4f}, val_acc{val_acc:.4f})这里没有直接用torch.optim.lr_scheduler.CosineAnnealingLR而是用ReduceLROnPlateau因为垃圾分类数据集的收敛曲线不像标准分类任务那么平滑初始学习率偏大时验证损失会上下震荡plateau 模式会在连续 3 个 epoch 不下降时自动减半。训练时把best_model.pth和最后一次 epoch 的权重分开保存因为轻微过拟合状态下验证准确率最高这个权重用于后续推理往往比最终权重效果好。如果数据集里“纸箱”有一万张而“电池”只有两百张交叉熵损失会偏向多数类。常见的做法是把F.cross_entropy的weight参数设为各类样本数的倒数class_counts torch.tensor([...]) class_weights 1.0 / class_counts.float() class_weights class_weights / class_weights.sum() * num_classes criterion nn.CrossEntropyLoss(weightclass_weights)加了类别权重后多数类的梯度会被压低少数类的错误预测会被放大惩罚。代价是训练集 loss 偏大、验证准确率可能微降但混淆矩阵里少数类的召回率会明显提升。这个取舍在垃圾分类场景里通常是值得的因为有害垃圾和电池这类少数类才是分错成本最高的。4. 从混淆矩阵到 ONNX 导出识别性能的验证与部署形态4.1 只看准确率会漏掉的错误用混淆矩阵定位“类间混淆对”训练结束不要只记录 val_acc要单独跑一次完整验证集并画出混淆矩阵。垃圾分类项目里准确率 90% 和 92% 的模型差异往往不来自某一个类别的整体提升而是几对“易混淆类”的处理方式变了。打印出错误最多的前 10 对组合from sklearn.metrics import confusion_matrix import numpy as np all_preds, all_labels [], [] model.eval() with torch.no_grad(): for images, labels in val_loader: outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) np.fill_diagonal(cm, 0) sorted_pairs np.dstack(np.unravel_index(np.argsort(cm.ravel())[::-1], cm.shape))[0][:10] for pred_cls, true_cls in sorted_pairs: if cm[true_cls][pred_cls] 0: print(f真实类别 {classes[true_cls]} 被误判为 {classes[pred_cls]}: {cm[true_cls][pred_cls]} 次)np.fill_diagonal(cm, 0)把对角线清零后剩下的非零位置全部是错误配对。输出这些对之后去看对应图片的样本通常会发现两个规律要么是这两类垃圾在颜色或形状上高度相似比如干净的酸奶盒和白纸要么是数据集中这两类的拍摄背景重叠太多模型实际学的是背景特征而不是垃圾本身。后者的解决思路是把背景变化类似的两类样本重新整理或者增加针对性增强。4.2 模型压缩与边缘部署的一条快捷路径ONNX 导出 量化训练好的 PyTorch 权重不能直接放到边缘盒子上跑常见部署链路是先转 ONNX再根据目标设备选择精度格式。import torch.onnx model.load_state_dict(torch.load(./best_model.pth)) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, ./garbage_model.onnx, input_names[input], output_names[logits], dynamic_axes{input: {0: batch_size}}, opset_version11 )导出后检查 ONNX 文件的输入输出维度确认最后一层输出是类别数而不是 1000。很多压缩包项目会在导出这一步出错原因是用torchvision.models.resnet18()加载了未修改分类头的预训练模型模型没有换成num_classes的分类层。导出改成动态 batch 尺寸后同一个模型可以适配摄像头单帧输入和批量评测。量化方面分类模型的常用策略压缩方式精度影响推理加速部署方式动态量化INT8分类任务约 0.5% 损失CPU 提升 1.52 倍torch.quantization.quantize_dynamic静态量化INT8需校准集损失可控制在 1% 内CPU 提升 23 倍需要少量真实场景图片做校准ONNX Runtime FP16几乎无损失GPU 提升明显适合 Jetson 系列垃圾分类部署最常见的两个场景是树莓派/工控机上的 CPU 推理和 Jetson Nano 上的轻量 GPU 推理。前者用动态量化就够后者导出 FP16 ONNX 跑 TensorRT 更合适。5. 垃圾分类模型训练完最值得记下的三个工程细节5.1 类别映射表 JSON 化部署时最容易被忽视的坑训练用的class_to_idx只在训练代码里存在部署时常常被丢掉。离线推理时如果手里只有 ONNX 模型和 numpy 数组没有类别名输出就是一个数字。把类别映射在训练最后一步单独保存和权重文件放一起import json with open(./class_map.json, w) as f: json.dump(full_ds.classes, f, ensure_asciiFalse)这里直接用列表顺序对应索引即可json.load读出来也是有序的保证推理代码里pred_cls class_names[pred_idx]不会错位。这个 JSON 文件和best_model.pth、ONNX 文件建议打成同一个压缩包避免分发时漏文件。5.2 早停策略的两个参数patience 和 min_delta垃圾分类数据集的验证准确率通常在 epoch 10 到 15 之间进入平台期继续训练只会过拟合。早停的参数设置经验是 patience 取 5 到 7min_delta 取 0.001。min_delta 的作用是忽略精度波动比如验证 loss 从 1.23 到 1.22 这种变化不算“变好”避免模型保存过于频繁。配合ReduceLROnPlateau一起用时先降学习率再早停比直接停的效果好。best_val_loss float(inf) epochs_no_improve 0 patience 5 for epoch in range(30): val_loss validate(model, val_loader) if val_loss best_val_loss - 0.001: best_val_loss val_loss epochs_no_improve 0 torch.save(model.state_dict(), ./best_model.pth) else: epochs_no_improve 1 if epochs_no_improve patience: print(fEarly stopping at epoch {epoch}) break5.3 测试时增强验证集上多跑几次取平均预测TTA 是把同一张图片做多次不同预处理后取平均输出。对单张图片分别做 CenterCrop 和四角裁剪五张图的结果做 softmax 平均可以在不训练的情况下把准确率提升 0.5 到 1 个百分点。实现时注意 batch 维度堆叠把五次增强结果一起 forward然后torch.mean(torch.softmax(logits, dim1), dim0)取平均。测试时增强会多花五倍推理时间适合对单张照片精度有要求的场景实时视频流里基本不用。本文还有配套的精品资源点击获取