基于ResNet50的垃圾识别分类实战:从数据清洗到模型部署

发布时间:2026/9/16 2:54:13
基于ResNet50的垃圾识别分类实战:从数据清洗到模型部署 简介这份压缩包是一套基于深度学习的垃圾识别分类完整工程面向掌握Python基础、希望将卷积神经网络应用于图像分类的开发者、学生或环保项目实践者。项目以可回收物、有害垃圾、湿垃圾、干垃圾等分类为目标完整覆盖数据预处理、图像增强、预训练模型微调、损失函数与优化器设置、训练监控及测试集评估等关键环节可支撑课程设计、毕业设计或垃圾分类应用原型搭建。压缩包共191个文件主体为176张标注图片与5个Python训练/推理脚本另含2个pth预训练权重、2个pyc缓存文件、4个xml配置及json等环境描述文件整体仅21.76MB轻量便于下载与本地复现。目前已有115人学习下载读者可获得可直接运行的代码、样本数据与权重文件快速体验从数据准备到模型评估的完整流程并可根据需求替换数据集或调整网络结构进行二次开发。1. 垃圾识别分类的真实起点类别体系先于网络结构同样的 ResNet50有人在公开数据集上跑到 92%你把实地拍摄的垃圾照片丢进去只有 74%问题多半不在网络上而在类别体系和数据分布。垃圾识别分类看起来是图像分类的入门题目实际是细粒度识别纸碗和纸杯、铝罐和铁罐人眼都要犹豫两秒。这个项目的完整链路是先定分类体系再做数据清洗和增强换预训练网络迁移学习用混淆矩阵找误判最后把模型、标签映射和推理脚本一起打包成 zip 交付。适合课程设计、毕业设计也适合小团队接识别 demo。动手前把这条链路上的参数和坑过一遍能省一半以上的返工时间。2. 垃圾识别分类的数据集构建脏数据清理、增强与类别平衡参数2.1 先定类别体系按目录组织数据比写标注文件更稳垃圾识别分类的第一步不是选网络而是把“到底分几类”定死。公开数据集常见做法是四分类或六分类四分类是厨余、可回收、有害、其他六分类会把可回收拆成纸类、塑料、玻璃、金属。类别粒度越细对数据量和标注质量的要求越高。我接这个类型项目时一般先按八二划分出 train 和 val每类训练图片不少于 500 张再谈模型。类别体系定下来后最省事的组织方式是目录即标签PyTorch 的torchvision.datasets.ImageFolder可以直接读这种结构data/ ├── train/ │ ├── glass/ │ ├── metal/ │ ├── paper/ │ ├── plastic/ │ └── food_waste/ └── val/ ├── glass/ ├── metal/ ├── paper/ ├── plastic/ └── food_waste/这里有一个容易踩的坑ImageFolder默认按字符序生成类别索引也就是 glass0、metal1、paper2、plastic3、food_waste4中文名按拼音排。我在训练前会单独导出一份label_map.json把索引和类别名固化成映射并跟训练脚本里的classes列表逐一比对。索引顺序一旦乱了后面 5.2 节的推理脚本输出就会张冠李戴而且这类错误在准确率上完全看不出来。目录整理好之后先跑一段统计脚本看各类数量分布这本节的平衡策略就基于这个分布来定。2.2 用 PIL verify 清理坏图和过小样本网络时代扒下来的垃圾图片经常混入表情包、截屏、损坏文件。一张坏图在DataLoader里会让整个 batch 训练中断这是新手最容易卡住的地方。我一般会在训练前做一次全量体检把打不开的、尺寸过小的文件移到bad_imgs备份目录而不是直接删from PIL import Image from pathlib import Path import shutil src Path(data/train) backup Path(data/bad_imgs) backup.mkdir(exist_okTrue) for img_path in src.rglob(*.*): if img_path.suffix.lower() not in {.jpg, .jpeg, .png, .bmp}: continue try: with Image.open(img_path) as im: im.verify() w, h im.size if w 96 or h 96: shutil.move(str(img_path), backup / img_path.name) except Exception: shutil.move(str(img_path), backup / img_path.name)这段代码里im.verify()只校验文件头部和内部结构不会完整解码像素跑几千张图也很快。异常路径和过小路径都做移动操作万一误判还能从备份里捞回来。96 像素这个阈值是经验值垃圾图片经过后续RandomResizedCrop增强后如果原图最短边低于 96裁剪到模型输入尺寸时目标特征已经严重失真。实际项目里如果发现坏图比例超过 5%先回去查采集链路而不是靠训练硬扛。2.3 增强参数旋转、色彩抖动与裁剪尺度垃圾照片的拍摄环境和训练集差异很大增强是低成本抗过拟合手段。对刚体物品小角度旋转比大幅仿射变换更安全色彩抖动则直接对标不同光照下的拍照效果。我在 torchvision 里常用这套配置from torchvision import transforms train_tf transforms.Compose([ transforms.RandomResizedCrop( size(224, 224), scale(0.6, 1.0), ratio(0.75, 1.33), ), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees(-15, 15), fill0), transforms.ColorJitter( brightness0.3, contrast0.3, saturation0.3, hue0.05, ), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ), ])RandomResizedCrop的scale(0.6, 1.0)保证至少保留原图 60% 面积避免把垃圾的小把手、瓶口这类判别区域裁掉ratio放宽到 0.75 到 1.33 模拟拍摄视角变化。RandomRotation只给 ±15 度因为垃圾识别是刚体分类倒扣的碗和正放的碗都是“碗”旋转 90 度这类极端增强反而会让模型学到错误的方向先验。ColorJitter的 hue 只能给 0.05给太大塑料和纸张的颜色边界会被打乱。这套增强在迁移学习里有一个前提Normalize的均值和标准差必须用 ImageNet 的默认值与预训练权重对齐不能自己拿数据集重新算否则预训练特征分布直接被破坏。验证集不要用ColorJitter和RandomRotation只做Resize、中心裁剪和归一化否则验证指标虚高。增强方法常见参数区间在垃圾识别里的作用RandomResizedCropscale0.6~1.0模拟拍摄距离和取景差异保留目标主体RandomRotationdegrees-15~15模拟随手丢放的任意角度幅度过大失真ColorJitterbrightness/contrast/saturation 0.2~0.4适配室内外光线、阴影、荧光灯做白平衡偏差RandomHorizontalFlipp0.5对称物体无副作用文字面上下翻转即可2.4 类别不平衡加权采样器的两个参数真实垃圾场景里塑料瓶和食品包装的数量远多于玻璃瓶要是直接按顺序采样模型会偏向多数类。WeightedRandomSampler是最直接的办法权重取类别样本数的倒数from torch.utils.data import WeightedRandomSampler class_count {0: 3200, 1: 1700, 2: 590, 3: 1200, 4: 3000} samples list(range(sum(class_count.values()))) weights [] for i in samples: cls sample_to_class[i] weights.append(1.0 / class_count[cls]) sampler WeightedRandomSampler( weights, num_sampleslen(samples), replacementTrue )num_samples一般设为整个数据集的长度配合replacementTrue每个 epoch 少数类会被重复抽样多数类被漏采的概率也随之升高。注意加权采样只解决采样频率不改变增强强度千万不要在加权采样之外再对少数类做暴力复制粘贴那样不仅过拟合还会让验证集里的同类样本显得“太好分类”。少数类样本少到每类不足 200 张时先考虑补充采集而不是一味加大采样权重。提示加了加权采样后训练 loss 曲线会比普通采样波动更大这是正常的。判断标准是验证集上每个类别的 recall 是否都起来了而不是只看整体准确率。3. 深度学习模型选型与两阶段训练ResNet50 的关键参数清单3.1 选型逻辑预训练权重能省多少样本垃圾识别没有大规模专有预训练模型ImageNet 预训练提供的边缘、纹理、形状特征是通用的这些基础特征在垃圾图像上依然有效。从零训练在单卡、几万张样本的规模下很难收敛所以迁移学习是这个任务的标准做法。选骨架时我按三个约束权衡训练显存、推理延迟、部署硬件。参数规模越大对细粒度纹理的区分能力越强但垃圾识别并非越强越好过度参数化在小数据集上反而掉点。模型参数量级适用场景部署注意ResNet18约 11MCPU 推理、小数据集材质相近的类别塑料 vs 纸容易混淆ResNet50约 25MGPU 训练、一般项目默认选型综合效果稳定TorchScript 导出顺畅EfficientNet-B0约 5M移动端或算力受限环境注意力相关算子需要较新版本 torchvision以 ResNet50 为例子分类头输出维度必须改成自己的类别数。这一步遗漏的话PyTorch 会在加载预训练权重时因为fc层维度不匹配直接报错所以也是很多人第一次跑迁移学习卡住的地方import torch.nn as nn model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) num_features model.fc.in_features model.fc nn.Linear(num_features, num_classes)in_features读的是原分类头输入维度ResNet50 是 2048。替换后的新fc层是随机初始化的这正是两阶段训练要解决的核心矛盾随机初始化的分类头在最初几个 epoch 会产生较大梯度如果整个网络一起训预训练骨干会被冲乱。3.2 两阶段训练参数冻结、解冻与早停我把训练拆成两个阶段。第一阶段冻结骨干网络只训练分类头让新分类头先适应预训练特征空间的分布第二阶段解冻骨干用更低的学习率微调整个网络。分类头训练阶段典型参数是lr5e-3微调阶段骨干学习率1e-5、分类头学习率1e-4。数值上分类头和骨干差 10 倍是常见做法因为骨干已经有成熟特征更新幅度必须更小。冻结操作和训练循环放在一起看更直观for param in model.backbone.parameters(): param.requires_grad False from torch.optim import AdamW from torch.optim.lr_scheduler import OneCycleLR optimizer AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr5e-3, weight_decay0.05, )filter只把requires_gradTrue的参数传给优化器冻结部分的骨干虽然有梯度存储空间但不会被更新省显存也避免误更新。lr5e-3针对的是分类头如果一上来就这么大权重初始化的随机噪声会把 pre 特征输出压住需要先加一点梯度裁剪或把 lr 降到1e-3跑几个 step 试试。阶段二解冻时分类头继续用相对大的1e-4骨干用1e-5。epoch 数量不是越多越好我用早停来控制best_acc, bad_epochs 0.0, 0 for epoch in range(max_epoch): train_one_epoch(model, train_loader, criterion, optimizer, scheduler) acc evaluate(model, val_loader) if acc best_acc: best_acc acc torch.save(model.state_dict(), best.pt) bad_epochs 0 else: bad_epochs 1 if bad_epochs 3: print(fearly stop at epoch {epoch}) breakpatience3表示验证集准确率连续 3 个 epoch 不创新高就停。很多项目跑到 20 个 epoch 结果过拟合早停能省一半时间。保存的是state_dict而不是整个模型因为后面第 5 章还要做 TorchScript 导出到时从脚本重新构建网络结构再加载权重这样管线更干净。3.3 优化器、warmup 和 OneCycle 参数数据集规模几万张以内时我优先用 AdamW 而不是 SGD。AdamW 对学习率不像 SGD 那么敏感配合weight_decay0.05能压住小样本下的过拟合SGD 加动量 0.9 和 nesterov 也可以但要额外调 warmup不然前几个 epoch 因为分类头随机初始化loss 会抖得很厉害。偏好 AdamW 的配置如下total_steps len(train_loader) * num_epochs scheduler OneCycleLR( optimizer, max_lr5e-4, total_stepstotal_steps, pct_start0.1, anneal_strategycos, )pct_start0.1表示前 10% 的步数学习率线性从低到高爬到max_lr这一步起到 warmup 作用之后按余弦曲线退火到接近 0。max_lr是阶段二解冻后的峰值5e-4在 batch size 32 时比较稳。如果把 batch size 翻倍到 64学习率也要跟着翻倍线性缩放规则在垃圾识别这类任务里实测直接可用。阶段一只用普通StepLR或干脆不接 scheduler因为分类头随机初始化阶段过早退火反而不利。4. 用混淆矩阵定位垃圾识别分类的误判并加置信度后处理4.1 从验证集生成分类报告和归一化混淆矩阵训练结束只看验证集准确率等于只看平均成绩掩盖了少数类的问题。垃圾识别分类里塑料和纸张被混淆、玻璃反光被误判成金属这些都要靠逐类的 precision、recall 和混淆矩阵暴露出来。评估脚本我放在evaluate.py里import torch import seaborn as sns import matplotlib.pyplot as plt from sklearn.metrics import classification_report, confusion_matrix preds, labels [], [] model.eval() with torch.no_grad(): for x, y in val_loader: out model(x) preds out.argmax(dim1).tolist() labels y.tolist() print(classification_report(labels, preds, target_namesclasses)) cm confusion_matrix(labels, preds) cm_norm cm / cm.sum(axis1, keepdimsTrue) sns.heatmap(cm_norm, annotTrue, fmt.2f, xticklabelsclasses, yticklabelsclasses) plt.savefig(confusion.png, dpi150)cm.sum(axis1, keepdimsTrue)是对每一行做归一化cm_norm[i][j]表示第 i 类的样本被预测成第 j 类的比例反映的是“召回视角”。fmt.2f保证热力图上显示小数而不是大整数否则样本量不均时满屏都是 500看不出比例。拿到报告后我先看三件事哪一类 recall 低于 0.8哪两类互相混整体准确率和 macro-F1 差多少。垃圾数据里 macro-F1 和 accuracy 差值超过 5 个百分点说明类别不平衡在起作用需要回到第 2 章的采样器调权重。4.2 三类典型误判及其修正方向热力图上最显眼的就是主对角线以外的亮块。我把垃圾识别里的常见误判分成三类按频率排序处理误判现象常见原因优先修正手段纸碗/纸杯、塑料盒/餐盒这类相似材质混淆形状和颜色接近模型主要靠颜色而非纹理提高输入分辨率到 320或在增强里增加灰度化分支金属反光被识别成塑料/纸张高光区域把特征冲淡训练集中反光样本少增强里加重 brightness 和 contrast补充反光样本遮挡或截断的垃圾如垃圾桶口的塑料袋训练集都是完整摆拍图增强里把 RandomResizedCrop 的 scale 下限降到 0.4处理顺序有讲究先确认“模型哪些类别分不开”再回看这些类别在验证集里的原始图片看是人眼都能区分的还是人眼都存疑的。人眼都分不清的对属于类别定义问题回到第 2 章合并类目或约束采集规范人眼能分清模型分不清才轮到调增强、加分辨率、换模型骨架。4.3 阈值之后低置信度样本交给人工而不是硬判即使混淆矩阵修好了垃圾识别落地时仍会有边缘样本。与其让模型硬判一个大概率错的类别不如加一个置信度阈值把低置信度样本挡下来probs torch.softmax(logits, dim1) score, idx probs.max(dim1) hard_mask score 0.60 hard_indices torch.nonzero(hard_mask).flatten()score是最大 softmax 概率idx是对应类别索引。hard_mask是布尔张量标记置信度低于 0.60 的样本这些样本不进入自动分拣流程而是转人工复核队列。阈值 0.60 不是拍脑袋定的我用验证集里每类的混淆情况反推先看哪一类的误判样本大多集中在哪个置信区间再取“误判率突然升高”的那个点作为阈值。调低阈值会漏判调高阈值会增加人工工作量这个权衡要在项目初始和客户对齐。置信度后处理还可以做一件事对历史预测日志做二次分析看哪些低置信度样本被修正后原本混淆的两个类别是否被区分开。这个统计结果直接作为第 6 章难例回流的输入。5. 打包可交付的 zipTorchScript 导出、推理脚本与目录校验5.1 为什么用 TorchScript 而不是裸 state_dict交付 zip 的最终目标是对方解压后直接能用不依赖你本地的 Python 工程结构。裸的state_dict只是权重字典对方如果没有同一份models.py定义网络结构torch.load会直接报错。TorchScript 把网络结构和权重序列化到同一个文件里对方只需要torch.jit.load就能加载。垃圾识别分类的forward是线性的没有动态控制流torch.jit.trace即可import json import torch model.eval() dummy torch.randn(1, 3, 224, 224) traced_model torch.jit.trace(model, dummy) traced_model.save(garbage_cls.pt) labels { str(i): name for i, name in enumerate(classes) } with open(label_map.json, w, encodingutf-8) as f: json.dump(labels, f, ensure_asciiFalse, indent2)torch.jit.trace拿着一个固定大小的输入沿 forward 走一遍记录实际的张量流和算子序列。dummy的尺寸必须是(1, 3, 224, 224)和训练时的输入尺寸一致。label_map.json的 key 用字符串而不是整数因为 JSON 规范不支持整数 key这一步顺序必须与训练时的classes列表完全一致否则推理结果对不上。5.2 目录结构和推理脚本交付 zip 我习惯按这个目录组织解压后任何一台带 Python 的机器都能跑garbage_cls_pkg/ ├── garbage_cls.pt ├── label_map.json ├── classify_one.py ├── requirements.txt └── README.mdclassify_one.py是单文件推理入口不依赖训练代码库里的任何模块import argparse import json import torch from PIL import Image from torchvision import transforms parser argparse.ArgumentParser() parser.add_argument(--image, requiredTrue) parser.add_argument(--model, defaultgarbage_cls.pt) parser.add_argument(--label_map, defaultlabel_map.json) args parser.parse_args() tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225], ), ]) model torch.jit.load(args.model, map_locationcpu) model.eval() with open(args.label_map, encodingutf-8) as f: labels json.load(f) with torch.no_grad(): img Image.open(args.image).convert(RGB) out model(tf(img).unsqueeze(0)) prob torch.softmax(out, dim1)[0] score, idx prob.max(dim0) print(json.dumps({ class: labels[str(idx.item())], score: round(score.item(), 4), }, ensure_asciiFalse))推理脚本最容易被忽略的是预处理一致性Resize((224, 224))和训练时的RandomResizedCrop并不一致。训练时网络见过的是随机裁剪的局部特征推理时直接整图压缩成 224目标尺度和内容分布都会偏移。更稳的做法是Resize(256)后接CenterCrop(224)模拟训练里最常见的那次裁剪。unsqueeze(0)把单张图变成(1, 3, 224, 224)的 batch 维度map_locationcpu保证在服务器或电脑上都能无显卡运行。requirements.txt只写一行说明就好“torch 和 torchvision 版本保持与训练环境一致”把确切版本号交给 README 记录。这样避免给出一份对方装不上的版本组合。5.3 打包命令与交付前校验目录和脚本都就绪后用系统 zip 命令打包cd /path/to/ zip -r garbage_cls_pkg.zip garbage_cls_pkg/ unzip -t garbage_cls_pkg.zipunzip -t只做完整性测试不实际解压。交付前还要做四件事我用表格列出来每一项都实际跑一遍再发出去校验项方法通过标准模型可加载torch.jit.load(garbage_cls.pt)无报错、无缺算子警告标签对齐对比label_map.json与训练时 classes 顺序每个索引对应类别一致预处理一致用验证集一张图跑推理和训练管线输出对比输出类别和置信度基本一致压缩包完整unzip -t或 Pythonzipfile.testzip()返回 None注意如果对方反馈“模型能加载但结果不对”90% 是预处理不一致优先查推理脚本里的Resize、CenterCrop和Normalize三个参数。6. 让垃圾识别分类持续变准低置信度样本的自动归档与再训练6.1 把部署环境的误判变成下一轮训练集模型交付不是终点垃圾识别分类上线后拍摄环境会不断带来新样本。我保留一份推理日志每行一个 JSON包含时间、图片路径、预测类别、置信度。然后用一段小脚本把低置信度样本自动归档while read -r line; do score$(echo $line | jq -r .score) if [ $score -lt 0.60 ]; then mkdir -p hard_examples cp $(echo $line | jq -r .image) hard_examples/ fi done inference.log这里的 0.60 阈值与第 4 章置信度后处理保持一致。归档的图片不等于训练样本还要经过人工复核这一关把同一类别、拍摄角度、光照环境接近的难例挑出来合并到data/train_plus目录和原始训练集一起构成下一轮数据。人工复核成本高所以只处理低置信度样本那些 high-confidence 且预测正确的样本不需要进回流池。难例积累到每类 200 张左右就开始微调而不是重新训练整个模型python train.py --resume best.pt --lr 1e-5 --epochs 5 \ --data data/train_pluslr1e-5是第 3 章微调阶段骨干学习率的量级回流样本数量少学习率再大会把已收敛特征冲乱。epoch 控制在 5 个以内配合早停逻辑验证集指标不再上升就立刻停。6.2 回流时控制难例比例别让模型记题回流最容易犯的错是把 200 张难例算作一个重头样本反复翻来覆去地训结果模型在验证集上刷高遇到新场景又打回原形。微调时难例占训练集比例不要超过 20%原始样本必须保留否则模型会把难例的特点当成通用规律。更稳妥的变体是让难例参与加权采样但不提高采样倍数只让它们自然地被更多 epoch 看到。回流前后的对比方式用第 4 章的混淆矩阵脚本比较同一个测试集上的逐类 recall。如果某个被混淆的类别对有了明显改善说明难例确实命中短板如果整体没动去检查归档阈值是不是设高了低置信度样本里混入了太多类别边界本来就模糊的图。如果加了难例之后外卖餐盒那一类的 recall 没有回升下次采集就往光线不足的档口多拍照片。本文还有配套的精品资源点击获取