简介这份资源面向图像分类入门者与深度学习实践者围绕垃圾分类场景提供一套可直接运行的神经网络模板。包内包含完整的垃圾分类数据集、数据集制作流程、模型训练与预测脚本train.py负责训练predict.py完成推理并支持在图片上以中文标注干垃圾、湿垃圾、可回收垃圾、有害垃圾四类结果适合作为图像分类任务的参考范例。资源共1046个文件以1041张jpg图片构成训练与测试数据另有2个py脚本、1个h5模型文件、1个txt说明及1个mp4演示视频压缩包约824.68MB。目前已有1332人学习下载。借助这套材料读者可以理解从数据整理、模型搭建到预测输出的完整链路并基于TensorFlow与OpenCV快速迁移到其他图像分类项目节省环境配置与代码调试时间。1. 垃圾分类图像识别从一张照片到四个桶中间到底要过几道关你拍一张外卖餐盒的照片扔进模型它告诉你「其他垃圾」——这个动作背后其实串起了图像处理、图像分类和垃圾分类三个环节。图像处理负责把原始像素变成模型能吃的干净输入图像分类负责从特征里判断类别垃圾分类则是把分类结果映射到具体的投放规则。很多人一上来就找最新的图像分类模型直接拿预训练权重跑微调结果发现准确率卡在 70% 上不去回头一查训练集里「可回收物」和「其他垃圾」的图片本身就标得含糊。这个方向适合有基本 Python 能力、想做一个能跑通的端到端项目的工程师也适合拿它当图像分类入门练手的人。它不要求你从头训一个大模型但要求你把数据管线、预处理和评估这三块做扎实。下面按「先搞清楚要做什么、再动手搭管线、最后避开那些让人翻车的坑」的顺序展开。2. 垃圾分类图像分类的任务拆解与数据准备2.1 先分清「图像处理」和「图像分类」在项目里各干什么图像处理在这个项目里不是指用 OpenCV 做几个滤镜就完事它承担的是把原始照片变成尺寸统一、光照一致、背景干扰被压掉的输入张量。常见做法是读入图片后先做短边缩放再中心裁剪然后归一化到 ImageNet 的均值和标准差。图像分类则是把处理后的张量送进网络输出每个类别的概率。垃圾分类是业务层的事——模型输出「厨余垃圾」这个标签业务代码再把它映射到「湿垃圾」或「易腐垃圾」的投放口。三者串起来才是一个能用的系统。很多人把这三层混在一起调模型准确率低就去换 backbone其实问题出在预处理阶段训练时用了随机裁剪增强推理时却直接 resize分布对不上精度自然掉。我一般会把预处理参数写成一个配置字典训练和推理共用同一份避免这种低级翻车。2.2 垃圾分类数据集怎么找、怎么标、怎么分公开的垃圾分类数据集常见的有 TrashNet、TACO 以及国内一些比赛放出的版本。TrashNet 只有 6 类、约 2500 张适合跑通流程但不适合追求高精度。TACO 的标注是 COCO 格式类别更细但需要自己转成分类任务用的文件夹结构。如果找不到合适的公开集自己拍也行但要注意每个类别至少 300 张起步且拍摄角度、光照、背景要有变化否则模型学到的只是背景颜色。标注环节最容易出问题的是边界类别。「沾了油渍的纸盒」算可回收还是其他垃圾不同人标法不一样。我的做法是先写一份标注规范把每个类别的判定规则用文字固定下来标注时遇到拿不准的就单独放一个「待定」文件夹最后统一裁决。划分训练集、验证集、测试集时按 7:1.5:1.5 的比例分层抽样保证每个类别在三个集合里的比例一致。import os import random import shutil from pathlib import Path def split_dataset(src_dir, dst_dir, ratios(0.7, 0.15, 0.15), seed42): src_dir: 每个子文件夹是一个类别里面是图片 dst_dir: 输出 train/val/test 三个子目录 ratios: 训练/验证/测试比例 random.seed(seed) src Path(src_dir) dst Path(dst_dir) classes [d.name for d in src.iterdir() if d.is_dir()] for cls in classes: imgs list((src / cls).glob(*.*)) random.shuffle(imgs) n len(imgs) n_train int(n * ratios[0]) n_val int(n * ratios[1]) splits { train: imgs[:n_train], val: imgs[n_train:n_train n_val], test: imgs[n_train n_val:] } for split, files in splits.items(): out dst / split / cls out.mkdir(parentsTrue, exist_okTrue) for f in files: shutil.copy2(f, out / f.name) print(f划分完成类别数{len(classes)}) split_dataset(raw_dataset, dataset)这段脚本按类别分层抽样保证每个类别在三个集合里的比例一致。ratios参数可以根据数据量调整数据少的时候验证集和测试集可以各留 10%。seed固定住是为了让每次划分结果可复现方便对比不同模型的实验。注意shutil.copy2会保留文件元信息如果磁盘空间紧张可以改成shutil.move。2.3 用 torchvision 搭一条可复用的预处理管线预处理管线要同时服务于训练和推理区别只在于训练时加随机增强、推理时用确定性变换。下面这条管线是我在垃圾分类项目里反复用过的版本输入尺寸 224训练时随机翻转加颜色抖动推理时只做 resize 和归一化。from torchvision import transforms # 训练和推理共用的归一化参数 IMAGENET_MEAN [0.485, 0.456, 0.406] IMAGENET_STD [0.229, 0.224, 0.225] train_tf transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD), ]) val_tf transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD), ])RandomResizedCrop的scale参数控制裁剪区域占原图的比例垃圾分类图片里主体通常占比较大所以下限设 0.7 而不是默认的 0.08避免裁到只剩背景。ColorJitter的强度不要开太大垃圾分类里颜色是重要线索——比如绿色通常对应厨余抖动太强会把颜色信息破坏掉。推理时用Resize(256)加CenterCrop(224)是标准做法和训练时的随机裁剪形成互补。3. 选模型与训练从 ResNet 到 Transformer 的取舍3.1 垃圾分类场景下 backbone 怎么选垃圾分类的类别数通常在 4 到 10 之间数据量从几千到几万张不等。这个规模下ResNet-50 仍然是一个很难被替代的基线。它在 ImageNet 上预训练后冻结前面几层、只微调后面两个 stage就能在几千张图上跑到 85% 以上的准确率。如果数据量更少ResNet-18 或 EfficientNet-B0 更合适参数少、过拟合风险低。Transformer 类模型比如 ViT 和 Swin在数据量充足时上限更高但它们对数据增强和正则化更敏感。我试过在 TrashNet 上直接微调 ViT-Base准确率反而不如 ResNet-50原因是数据量不够注意力机制没学到有效的局部特征。如果要用 Transformer建议先用 ResNet 跑一个基线确认数据管线没问题之后再换。最新的图像分类模型不一定是你的最优解这一点在垃圾分类这种细粒度不高的任务上尤其明显。3.2 微调 ResNet-50 的完整训练脚本与参数说明下面这个训练脚本是我在多个垃圾分类项目里用过的模板核心逻辑是加载预训练权重、替换最后的全连接层、分层设置学习率、用余弦退火调度。import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, models from torch.optim.lr_scheduler import CosineAnnealingLR from torchvision.transforms import Compose def build_model(num_classes, freeze_backboneTrue): model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V2) if freeze_backbone: for name, param in model.named_parameters(): if layer4 not in name and fc not in name: param.requires_grad False model.fc nn.Linear(model.fc.in_features, num_classes) return model def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total 0.0, 0, 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * imgs.size(0) correct (outputs.argmax(1) labels).sum().item() total imgs.size(0) return total_loss / total, correct / total # 数据加载 train_ds datasets.ImageFolder(dataset/train, transformtrain_tf) val_ds datasets.ImageFolder(dataset/val, transformval_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) device torch.device(cuda if torch.cuda.is_available() else cpu) model build_model(num_classeslen(train_ds.classes)).to(device) # 分层学习率新加的全连接层用大学习率微调的 layer4 用小学习率 params [ {params: model.fc.parameters(), lr: 1e-3}, {params: model.layer4.parameters(), lr: 1e-4}, ] optimizer torch.optim.AdamW(params, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max30) criterion nn.CrossEntropyLoss(label_smoothing0.1) for epoch in range(30): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) scheduler.step() print(fEpoch {epoch1}: loss{train_loss:.4f}, acc{train_acc:.4f})build_model里冻结了除layer4和fc之外的所有层这是数据量在几千张时的稳妥做法。如果数据超过两万张可以把layer3也解冻。label_smoothing0.1是防止模型对某个类别过度自信垃圾分类里有些图片确实模棱两可硬标签会逼模型记住噪声。CosineAnnealingLR的T_max设成总 epoch 数让学习率从初始值平滑降到接近零。AdamW的weight_decay设 1e-4比 SGD 的 5e-4 更温和适合微调场景。3.3 训练过程中该盯哪些指标训练时不要只看 loss 和 accuracy。垃圾分类项目里我必看三个东西每个类别的召回率、混淆矩阵、以及验证集 loss 和训练集 loss 的差值。如果「可回收物」的召回率明显低于其他类别大概率是这一类里混进了太多「其他垃圾」的图片或者这一类内部差异太大——比如塑料瓶和纸箱都被标成可回收但视觉上差别很大。混淆矩阵能直接告诉你哪两个类别在互相误判。我遇到过一次「厨余垃圾」和「其他垃圾」互相误判严重查了半天发现是拍摄时厨余垃圾经常用黑色塑料袋装着而其他垃圾里也有大量黑色塑料袋模型学到的其实是「黑色塑料袋」这个背景特征。解决办法是在预处理阶段加一个简单的背景裁剪或者补充一些散装厨余垃圾的图片。4. 推理部署与数据增强的实战细节4.1 把训练好的模型导出成推理脚本训练完的模型要能脱离训练代码独立跑。下面这个推理脚本加载 checkpoint对单张图片输出类别和置信度适合集成到 Web 服务或边缘设备上。import torch from PIL import Image from torchvision import transforms def load_model(ckpt_path, num_classes, device): model models.resnet50(weightsNone) model.fc nn.Linear(model.fc.in_features, num_classes) model.load_state_dict(torch.load(ckpt_path, map_locationdevice)) model.to(device).eval() return model def predict(image_path, model, class_names, device): img Image.open(image_path).convert(RGB) tensor val_tf(img).unsqueeze(0).to(device) with torch.no_grad(): logits model(tensor) probs torch.softmax(logits, dim1) conf, idx probs.max(1) return class_names[idx.item()], conf.item() model load_model(best.pth, num_classes6, devicedevice) label, score predict(test.jpg, model, train_ds.classes, device) print(f预测{label}置信度{score:.3f})推理时一定要用model.eval()和torch.no_grad()前者关闭 dropout 和 batch norm 的训练行为后者省显存。val_tf必须和训练时的验证变换完全一致包括 resize 的尺寸和归一化参数。如果部署环境没有 GPUmap_location要设成cpu同时把模型转成 half 精度可以进一步压缩体积。4.2 数据增强在垃圾分类里的边界数据增强不是越多越好。垃圾分类任务里水平翻转是安全的因为垃圾不会因为左右翻转改变类别。垂直翻转要谨慎——有些类别的图片上下颠倒后看起来像另一个类别。旋转角度超过 15 度就可能把「杯子」转成「碗」的视角。颜色抖动前面说过强度要控制。我一般会先用一组保守的增强跑基线然后逐个加入更强的增强看验证集准确率的变化。如果加入某个增强后验证集掉了两个点以上就说明这个增强和任务语义冲突。MixUp 和 CutMix 在垃圾分类上效果不稳定因为混合后的图片可能同时包含两个类别的特征而垃圾分类的标签是互斥的。如果要用建议只在数据量很大且类别边界清晰时尝试。4.3 用混淆矩阵定位系统性误判混淆矩阵不只是一个评估工具它可以直接指导你补数据。下面这段代码画出归一化的混淆矩阵并找出误判最多的类别对。import numpy as np import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix def plot_confusion(model, loader, class_names, device): model.eval() all_preds, all_labels [], [] with torch.no_grad(): for imgs, labels in loader: imgs imgs.to(device) preds model(imgs).argmax(1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds, normalizetrue) fig, ax plt.subplots(figsize(8, 6)) im ax.imshow(cm, cmapBlues) ax.set_xticks(range(len(class_names))) ax.set_yticks(range(len(class_names))) ax.set_xticklabels(class_names, rotation45, haright) ax.set_yticklabels(class_names) for i in range(len(class_names)): for j in range(len(class_names)): ax.text(j, i, f{cm[i, j]:.2f}, hacenter, vacenter) plt.colorbar(im) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150) plot_confusion(model, val_loader, train_ds.classes, device)归一化后的混淆矩阵每一行加起来是 1对角线上的值就是该类别的召回率。如果第 i 行第 j 列的值超过 0.2说明类别 i 有超过 20% 的样本被误判成类别 j这时候就要去翻这些误判样本的原图看是标注问题还是特征混淆。我一般会把误判样本的路径导出来人工过一遍确认是标注错误就修正标签确认是特征混淆就补拍类似场景的图片。5. 避坑与排查垃圾分类图像分类里最容易翻车的五件事5.1 训练集准确率 99% 但测试集只有 60%现象训练 loss 一路降到 0.01训练准确率接近满分但验证集和测试集准确率卡在 60% 左右不动。原因最常见的是数据泄漏——同一张图片的副本同时出现在训练集和验证集里。垃圾分类数据经常从网上批量下载同一张图可能有多个文件名。另一个原因是训练集和测试集的拍摄条件差异太大比如训练集全是白底商品图测试集是手机随手拍。解决划分数据集之前先做去重用图片的 MD5 或感知哈希pHash找出重复项。如果训练集和测试集来源不同要么统一来源要么在训练时加入更强的风格增强比如随机调整亮度、对比度和背景模糊。5.2 模型把「黑色塑料袋」当成「其他垃圾」的唯一特征现象所有装在黑色塑料袋里的垃圾都被判成其他垃圾哪怕里面装的是塑料瓶。原因训练数据里「其他垃圾」类别大量使用黑色塑料袋拍摄模型学到了这个捷径特征而不是垃圾本身的形状和材质。解决补充散装垃圾的图片或者在预处理阶段做前景分割把背景裁掉。如果不想引入分割模型可以在训练时随机裁剪掉图片边缘区域逼模型关注中心主体。更彻底的做法是收集一批「黑色塑料袋装可回收物」的图片单独标成可回收让模型学会区分内容和容器。5.3 推理速度在 CPU 上慢到无法接受现象模型在 GPU 上单张推理 10ms部署到 CPU 服务器后变成 500msQPS 上不去。原因ResNet-50 在 CPU 上的浮点运算量不小加上 Python 的 PIL 解码和预处理也占时间。解决先把模型转成 ONNX 或 TorchScript再用 ONNX Runtime 或 OpenVINO 做推理加速。预处理阶段把 PIL 换成 OpenCV 的imdecode速度能快一倍。如果还不行换 MobileNetV3 或 EfficientNet-Lite精度掉两三个点但速度提升明显。批量推理时把 batch size 调到 8 或 16CPU 的利用率会高很多。5.4 新增一个类别后模型完全失效现象原本 6 类的模型要扩展到 7 类重新训练后所有类别的准确率都掉了。原因直接修改全连接层后新加的类别随机初始化训练时梯度会冲击已经学好的特征。如果学习率没调小整个网络都会被带偏。解决扩展类别时先把 backbone 全部冻结只训练新的全连接层 5 到 10 个 epoch等新类别有基本识别能力后再解冻 layer4 做微调。学习率要比第一次训练时小一个数量级。另外新类别的数据量不能太少至少要和最少的旧类别持平。5.5 验证集准确率波动大每次跑结果都不一样现象同样的代码和数据每次训练完验证集准确率在 82% 到 88% 之间跳。原因随机种子没固定数据划分、权重初始化、数据增强的随机性都会影响结果。另外 batch size 太小也会导致梯度噪声大。解决在脚本开头固定torch.manual_seed、np.random.seed和random.seedDataLoader 的worker_init_fn也要设种子。如果数据量允许把 batch size 提到 32 以上。评估时用 K 折交叉验证代替单次划分取平均准确率作为最终指标波动会小很多。6. 用 Grad-CAM 验证模型到底在看哪里模型告诉你「厨余垃圾」但它到底是看到了剩饭剩菜还是看到了你拍照时桌子的木纹这个问题不搞清楚上线之后迟早翻车。Grad-CAM 是我在垃圾分类项目里必做的一步验证它把模型最后一块卷积层的梯度加权回特征图生成一张热力图告诉你模型做判断时关注了图片的哪个区域。import cv2 import numpy as np import torch import torch.nn.functional as F class GradCAM: def __init__(self, model, target_layer): self.model model self.target_layer target_layer self.gradients None self.activations None target_layer.register_forward_hook(self._save_activation) target_layer.register_full_backward_hook(self._save_gradient) def _save_activation(self, module, input, output): self.activations output.detach() def _save_gradient(self, module, grad_input, grad_output): self.gradients grad_output[0].detach() def generate(self, input_tensor, class_idxNone): self.model.eval() output self.model(input_tensor) if class_idx is None: class_idx output.argmax(1).item() self.model.zero_grad() output[0, class_idx].backward() weights self.gradients.mean(dim(2, 3), keepdimTrue) cam (weights * self.activations).sum(dim1, keepdimTrue) cam F.relu(cam) cam F.interpolate(cam, sizeinput_tensor.shape[2:], modebilinear, align_cornersFalse) cam cam.squeeze().cpu().numpy() cam (cam - cam.min()) / (cam.max() - cam.min() 1e-8) return cam, class_idx # 使用示例 grad_cam GradCAM(model, model.layer4[-1]) img_tensor val_tf(Image.open(test.jpg).convert(RGB)).unsqueeze(0).to(device) cam, pred_idx grad_cam.generate(img_tensor) # 叠加到原图 img cv2.imread(test.jpg) img cv2.resize(img, (cam.shape[1], cam.shape[0])) heatmap cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET) overlay cv2.addWeighted(img, 0.6, heatmap, 0.4, 0) cv2.imwrite(gradcam_result.jpg, overlay) print(f预测类别{train_ds.classes[pred_idx]})target_layer选model.layer4[-1]是因为它是 ResNet 最后一个卷积块感受野足够大能覆盖整个垃圾主体。register_full_backward_hook拿到的梯度是反向传播时该层的梯度和激活值相乘再求和就得到每个通道的权重。F.relu过滤掉负值因为只有正贡献的区域才是模型真正关注的。最后归一化到 0 到 1 再叠加热力图。拿到热力图之后怎么看如果高亮区域集中在垃圾主体上说明模型学到了正确的特征如果高亮区域在背景、桌沿、或者图片角落说明模型走了捷径。我遇到过一次模型把「可回收物」的判断依据放在图片右下角的水印上因为训练集里可回收物的图片大多来自同一个网站水印位置固定。解决办法是把图片边缘裁掉 10%或者补充不同来源的图片。这个验证步骤花不了多少时间但能帮你提前发现那些准确率数字掩盖不了的问题。我现在的习惯是每训练完一个版本先跑一批测试图的 Grad-CAM确认关注区域合理之后再去看准确率报表。希望帮到你。本文还有配套的精品资源点击获取