简介这份深度学习图像分类数据集面向从事计算机视觉入门与垃圾分类识别实践的开发者、学生及算法爱好者围绕塑料瓶、玻璃瓶、金属瓶等可回收物类别构建可直接用于训练与评估卷积神经网络分类模型也可作为课程设计、毕业项目或算法对比实验的数据基础。资源包共约2000个文件以1998张jpg图像为主体另附1个py脚本与1个json类别文件压缩包大小约39.65MB图像按类别目录分门别类存放便于直接接入ImageFolder等标准加载方式。数据已预先划分为训练集约2200张、验证集约200张、测试集约100张省去自行切分的步骤json文件可用于核对具体类别名称与索引映射。目前已有269人学习下载适合希望快速验证分类网络效果、练习数据增强与迁移学习流程的读者参考使用。1. 垃圾分类图像分类数据集从「这堆图能不能用」到跑通第一个基线你手头可能有一批从社区回收站、校园垃圾桶旁、或者公开渠道攒下来的瓶子与垃圾照片几百到几千张不等文件夹里混着塑料瓶、玻璃瓶、纸盒、果皮命名乱七八糟。现在你想用深度学习做一个图像分类模型把「可回收瓶类」「其他垃圾」这些类别自动分开。这件事的核心不是模型多深而是数据集能不能撑起训练——类别是否均衡、标注是否一致、图像质量是否够用。垃圾分类、瓶子垃圾图像分类这个方向本质是一个细粒度、强背景干扰的视觉分类任务适合刚接触深度学习图像分类的从业者拿来练手也适合做环保、回收分拣场景的工程团队做原型验证。这一章先把「这是什么、能解决什么、适合谁」讲清楚后面几章再落到具体怎么做、参数怎么设、坑在哪。2. 垃圾分类数据集长什么样类别体系、标注格式与选型理由2.1 先定类别体系再谈图像数量很多人一上来就问「多少张图够」但更关键的问题是「分几类」。垃圾分类的类别体系直接决定模型能不能落地。常见做法是两级一级粗分可回收物、厨余垃圾、有害垃圾、其他垃圾二级细分塑料瓶、玻璃瓶、易拉罐、纸盒等。如果你只做瓶子垃圾图像分类建议先聚焦 4 到 6 个视觉差异明显的类比如「透明塑料瓶」「有色塑料瓶」「玻璃瓶」「易拉罐」「纸杯」「其他」。类别太少模型学不到区分度类别太多且视觉边界模糊比如「可回收塑料」和「其他塑料」会让标注一致性崩掉。我一般会先画一张类别定义表把每个类别的判定标准写清楚比如「透明塑料瓶瓶身透光、无标签或标签面积小于 30%」。这张表不是给模型看的是给标注的人看的。没有这张表后面标注出来的数据大概率是玄学。类别名判定要点易混淆项建议最少样本数透明塑料瓶瓶身透光标签面积小玻璃瓶、有色塑料瓶300有色塑料瓶瓶身有颜色不透光透明塑料瓶300玻璃瓶材质反光瓶壁厚透明塑料瓶250易拉罐金属反光圆柱形纸杯250纸杯/纸盒纸质纹理无金属反光易拉罐200其他垃圾不属于以上类别全部200这张表里的「建议最少样本数」不是硬性标准而是经验下限。如果某个类少于 200 张模型大概率会在这个类上翻车。类别不均衡时优先补少样本类而不是继续堆多样本类。2.2 标注格式文件夹分类 vs CSV 清单图像分类最省事的标注方式就是按文件夹分。每个类别一个文件夹文件夹名就是类别名图片直接放进去。这种格式几乎所有训练框架都认比如torchvision.datasets.ImageFolder、tensorflow.keras.preprocessing.image_dataset_from_directory。另一种是 CSV 清单两列filename,label。CSV 的好处是灵活可以随时改标签、做交叉验证划分也方便记录来源和备注。我一般会同时保留两份原始文件夹结构用于快速预览CSV 用于训练脚本读取。下面是一个把文件夹结构转成 CSV 的脚本顺便做一次基础校验。import os import csv from pathlib import Path # 数据集根目录下面每个子文件夹是一个类别 data_root Path(dataset/garbage_cls) # 输出 CSV 路径 csv_path Path(dataset/train_labels.csv) # 允许的图片后缀避免把 .DS_Store 之类的文件读进来 valid_ext {.jpg, .jpeg, .png, .bmp, .webp} rows [] class_names sorted([d.name for d in data_root.iterdir() if d.is_dir()]) print(类别列表:, class_names) for cls_name in class_names: cls_dir data_root / cls_name for img_path in cls_dir.iterdir(): if img_path.suffix.lower() not in valid_ext: continue # 相对路径方便后续迁移数据集位置 rel_path img_path.relative_to(data_root) rows.append({filename: str(rel_path), label: cls_name}) # 写 CSV带表头 with csv_path.open(w, newline, encodingutf-8) as f: writer csv.DictWriter(f, fieldnames[filename, label]) writer.writeheader() writer.writerows(rows) print(f共写入 {len(rows)} 条记录到 {csv_path})这段脚本的逻辑很直白遍历每个类别文件夹过滤掉非图片文件把相对路径和类别名写成 CSV。参数上data_root指向你的数据集根目录csv_path是输出位置。跑完之后你可以用pandas读一下检查每个类别的数量是否均衡。import pandas as pd df pd.read_csv(dataset/train_labels.csv) print(df[label].value_counts())如果某个类别的数量明显偏少比如只有 80 张而其他类有 500 张那就需要先处理不均衡问题而不是直接开训。常见做法是数据增强、重采样、或者用类别权重。这个后面在训练章节会展开。2.3 图像质量与背景干扰垃圾分类数据集的天然难点垃圾分类图像和 ImageNet 那种干净数据集不一样。你的图片可能来自手机拍摄背景里有地面、手、其他垃圾、光照变化。瓶子类还有反光、透明、形变的问题。透明塑料瓶和玻璃瓶在视觉上非常接近模型很容易把两者搞混。这不是模型不行是数据本身就有歧义。我一般会做三件事第一人工抽检 100 张看看有没有明显标错的第二统计图像尺寸分布太小的图比如小于 100x100直接剔除第三对透明瓶和玻璃瓶这类易混类别额外补充侧光、逆光、不同角度的样本。如果条件允许用albumentations做在线增强时对这两类加更强的颜色抖动和随机裁剪让模型学到材质差异而不是背景捷径。提示不要用「瓶子」和「垃圾」这种粗标签直接训。粗标签会让模型学到「有瓶子就是可回收」这种捷径换个场景就废了。3. 用 PyTorch 跑通瓶子垃圾图像分类基线从 DataLoader 到第一个 epoch3.1 环境准备与最小依赖深度学习环境这件事新手最容易卡在版本冲突上。我一般用 conda 建一个干净环境只装必要的东西。下面这套是 CPU 和单卡 GPU 都能跑的配置不追求最新追求能复现。conda create -n garbage_cls python3.10 -y conda activate garbage_cls pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install pandas pillow scikit-learn matplotlib tqdm如果你没有 GPU把cu118换成cpu即可。装完之后用下面这行验证import torch print(torch.__version__, torch.cuda.is_available())torch.cuda.is_available()返回True说明 GPU 可用。返回False也不影响跑基线只是慢一点。3.2 构建 Dataset 和 DataLoaderPyTorch 里最稳的做法是继承torch.utils.data.Dataset自己写__getitem__。这样你能完全控制读取逻辑也方便后面加增强。import torch from torch.utils.data import Dataset, DataLoader from PIL import Image from torchvision import transforms import pandas as pd from pathlib import Path class GarbageDataset(Dataset): def __init__(self, csv_file, img_root, transformNone): self.df pd.read_csv(csv_file) self.img_root Path(img_root) self.transform transform # 类别到索引的映射排序保证可复现 self.classes sorted(self.df[label].unique()) self.class_to_idx {c: i for i, c in enumerate(self.classes)} def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] img_path self.img_root / row[filename] # 统一转 RGB防止灰度图或 RGBA 图导致通道数不一致 image Image.open(img_path).convert(RGB) label self.class_to_idx[row[label]] if self.transform: image self.transform(image) return image, label # 训练增强随机裁剪、翻转、颜色抖动 train_tf transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.3), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) # 验证集只做 resize 和归一化不做随机增强 val_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]), ]) train_ds GarbageDataset(dataset/train_labels.csv, dataset/garbage_cls, train_tf) val_ds GarbageDataset(dataset/val_labels.csv, dataset/garbage_cls, val_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4)这里有几个参数值得说清楚。batch_size32是单卡 8G 显存下的稳妥值显存不够就降到 16 或 8。num_workers4在 Windows 上有时会出问题如果报错就改成 0。Normalize用的均值和方差是 ImageNet 的统计值因为后面要用预训练模型保持一致最省事。RandomResizedCrop的scale(0.7, 1.0)表示随机裁取原图 70% 到 100% 的区域对瓶子这种主体居中的图比较合适裁太狠会把瓶身切掉。3.3 选一个 backboneResNet18 还是更小的模型垃圾分类数据集通常不大几千到几万张。这种规模下ResNet18 或 EfficientNet-B0 是性价比最高的选择。ResNet18 参数量约 1100 万ImageNet 预训练权重容易拿训练稳定。EfficientNet-B0 更小但输入分辨率要求 224 以上对透明瓶的细节可能更友好。我一般先用 ResNet18 跑基线如果准确率卡在某个值上不去再换 backbone 对比。import torch.nn as nn from torchvision import models def build_model(num_classes, backboneresnet18, pretrainedTrue): if backbone resnet18: model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT if pretrained else None) # 替换最后的全连接层输出类别数 model.fc nn.Linear(model.fc.in_features, num_classes) elif backbone efficientnet_b0: model models.efficientnet_b0(weightsmodels.EfficientNet_B0_Weights.DEFAULT if pretrained else None) model.classifier[1] nn.Linear(model.classifier[1].in_features, num_classes) else: raise ValueError(f不支持的 backbone: {backbone}) return model num_classes len(train_ds.classes) model build_model(num_classes, backboneresnet18, pretrainedTrue) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) print(f类别数: {num_classes}, 设备: {device})pretrainedTrue会加载 ImageNet 权重这对小数据集非常关键。从头训 ResNet18 在几千张图上大概率过拟合。替换fc层之后只有最后一层是随机初始化的前面都是预训练特征。3.4 训练循环与关键参数训练循环本身不复杂但有几个参数决定你能不能跑出结果。优化器用AdamW学习率1e-3对预训练模型微调偏大我一般用3e-4或1e-4。损失函数用CrossEntropyLoss如果类别不均衡加weight参数。import torch.optim as optim from tqdm import tqdm criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max20) def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total 0.0, 0, 0 for images, labels in tqdm(loader, desctrain): images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() total_loss, correct, total 0.0, 0, 0 for images, labels in tqdm(loader, descval): images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total best_acc 0.0 for epoch in range(20): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fepoch {epoch1}: train_loss{train_loss:.4f} train_acc{train_acc:.4f} fval_loss{val_loss:.4f} val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_garbage_cls.pth) print(f保存最佳模型val_acc{best_acc:.4f})CosineAnnealingLR的T_max20表示 20 个 epoch 内学习率按余弦曲线降到接近 0。如果你只跑 10 个 epoch就把T_max改成 10。weight_decay1e-4是 AdamW 的常规值能压一点过拟合。保存模型时只存state_dict不存整个模型对象这样加载时更灵活。跑完 20 个 epoch如果验证集准确率能到 85% 以上说明数据集和基线是通的。如果卡在 60% 左右大概率是类别混淆或标注问题不是模型问题。4. 垃圾分类图像分类的避坑与排查标注、不均衡、过拟合、透明瓶混淆4.1 标注不一致同一个瓶子被标成两个类现象训练 loss 震荡验证准确率上不去混淆矩阵里某两个类互相错得特别多。原因标注时没有统一标准比如有人把「透明塑料瓶」标成「玻璃瓶」有人反过来。解决先抽 200 张做一致性检查让两个人独立标同一批图算一下一致率。低于 90% 就重新定标准把易混类别的判定规则写得更细比如「瓶底有厚玻璃纹路算玻璃瓶瓶身有塑料接缝算塑料瓶」。4.2 类别不均衡少样本类被模型忽略现象总体准确率看着还行但少样本类的召回率极低模型几乎全预测成多样本类。原因CrossEntropyLoss默认每个样本权重一样多样本类主导梯度。解决给CrossEntropyLoss加weight权重和类别频率成反比。或者用重采样让每个 batch 里少样本类被抽到的概率更高。import numpy as np class_counts train_ds.df[label].value_counts().sort_index().values class_weights 1.0 / class_counts class_weights class_weights / class_weights.sum() * len(class_counts) class_weights torch.tensor(class_weights, dtypetorch.float32).to(device) criterion nn.CrossEntropyLoss(weightclass_weights) print(类别权重:, class_weights)这段代码先统计每个类的样本数取倒数再归一化到均值为 1 左右。weight传给损失函数后少样本类的 loss 会被放大梯度更新时更受重视。4.3 过拟合训练准确率 99%验证准确率 70%现象训练集准确率一路涨到 99%验证集卡在 70% 不动甚至下降。原因数据集太小、模型太大、增强不够。解决先加数据增强RandomResizedCrop、ColorJitter、RandomRotation都可以上。再加DropoutResNet18 的fc前面可以插一个nn.Dropout(0.5)。最后考虑冻结前面层只训fc或者用更小的学习率。4.4 透明瓶与玻璃瓶混淆模型学的是背景不是材质现象混淆矩阵里透明塑料瓶和玻璃瓶互相错得最多但人眼看这两类差异明显。原因模型可能学到了背景捷径比如透明瓶常出现在某类背景上玻璃瓶在另一类背景上。解决把背景裁掉或者用分割模型先抠出瓶身再分类。更简单的做法是加随机背景替换增强把瓶子贴到不同背景上逼模型看瓶身。4.5 图像尺寸和通道不一致导致训练中断现象训练到一半报错提示 tensor 尺寸不匹配或通道数不对。原因数据集里混了灰度图、RGBA 图、尺寸差异极大的图。解决在__getitem__里统一convert(RGB)在transform里统一Resize。如果还有问题写个脚本扫一遍所有图片把打不开的、尺寸小于 50x50 的直接列出来剔除。from PIL import Image from pathlib import Path bad_files [] for img_path in Path(dataset/garbage_cls).rglob(*): if img_path.suffix.lower() not in {.jpg, .jpeg, .png, .bmp, .webp}: continue try: with Image.open(img_path) as im: w, h im.size if w 50 or h 50: bad_files.append((str(img_path), too_small, w, h)) except Exception as e: bad_files.append((str(img_path), unreadable, str(e))) for item in bad_files: print(item) print(f共发现 {len(bad_files)} 个问题文件)这个脚本会列出所有打不开或太小的图片。跑一遍把问题文件处理掉能省掉很多训练中途报错的麻烦。5. 把基线推到可用混淆矩阵诊断、难例挖掘与一个提点技巧5.1 用混淆矩阵定位问题类别准确率只是一个数字真正有用的是混淆矩阵。下面这段代码在验证集上跑一遍画出每个类别的错分情况。from sklearn.metrics import confusion_matrix, classification_report import numpy as np torch.no_grad() def get_all_preds(model, loader, device): model.eval() all_preds, all_labels [], [] for images, labels in loader: images images.to(device) outputs model(images) preds outputs.argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) return np.array(all_labels), np.array(all_preds) labels, preds get_all_preds(model, val_loader, device) cm confusion_matrix(labels, preds) print(混淆矩阵:) print(cm) print(classification_report(labels, preds, target_namesval_ds.classes))classification_report会给出每个类的 precision、recall、f1-score。如果某个类的 recall 特别低说明这个类被大量错分。结合混淆矩阵你能看出它主要被错分成了哪个类。比如透明塑料瓶的 recall 只有 0.6且大部分错分到玻璃瓶那就针对这两类做难例挖掘。5.2 难例挖掘把错分样本挑出来单独看难例挖掘的思路很简单用当前模型跑一遍训练集把预测错误的样本挑出来人工看一遍确认是标注问题还是模型问题。如果是标注问题修正标签如果是模型问题把这些难例复制几份加入训练集或者对这些类做更强的增强。torch.no_grad() def find_hard_examples(model, dataset, device, top_k50): model.eval() hard [] for idx in range(len(dataset)): image, label dataset[idx] # dataset 返回的是增强后的图这里为了诊断用原始图更合适 # 简化处理直接用增强后的图推理 input_tensor image.unsqueeze(0).to(device) output model(input_tensor) pred output.argmax(dim1).item() if pred ! label: hard.append((idx, label, pred)) return hard[:top_k] hard_examples find_hard_examples(model, val_ds, device, top_k30) for idx, true_label, pred_label in hard_examples: print(fidx{idx} 真实{val_ds.classes[true_label]} 预测{val_ds.classes[pred_label]})这段代码会列出验证集里被错分的样本索引和类别。你可以根据索引找到原图看看是标注错了还是模型确实分不出来。如果是标注错了改标签如果是模型问题把这些图加入训练集或者对这类图做针对性增强。5.3 一个提点技巧两阶段训练与标签平滑如果基线准确率卡在 85% 左右上不去可以试两阶段训练。第一阶段冻结 backbone只训分类头学习率设1e-3跑 5 个 epoch。第二阶段解冻全部层学习率降到1e-4跑 15 个 epoch。这样能让分类头先适应数据分布再微调整个网络比一上来就全量微调更稳。另外标签平滑label smoothing对垃圾分类这种有歧义的类别很有效。CrossEntropyLoss(label_smoothing0.1)会让模型不要把概率压到 1.0对易混类别更宽容。我一般会在第二阶段加上通常能提 1 到 2 个点。# 第二阶段解冻全部层小学习率微调 for param in model.parameters(): param.requires_grad True optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) criterion nn.CrossEntropyLoss(label_smoothing0.1) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max15) for epoch in range(15): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fstage2 epoch {epoch1}: val_acc{val_acc:.4f})标签平滑的0.1是经验值太大比如 0.3会让模型欠拟合太小比如 0.01效果不明显。两阶段训练的关键是第二阶段学习率一定要小否则会破坏预训练特征。我自己的习惯是每次跑完基线先看混淆矩阵再决定是补数据、改标注还是调模型。数据集的问题永远比模型多。希望帮到你。本文还有配套的精品资源点击获取
垃圾分类图像分类实战：从数据集构建到PyTorch基线模型
简介这份深度学习图像分类数据集面向从事计算机视觉入门与垃圾分类识别实践的开发者、学生及算法爱好者围绕塑料瓶、玻璃瓶、金属瓶等可回收物类别构建可直接用于训练与评估卷积神经网络分类模型也可作为课程设计、毕业项目或算法对比实验的数据基础。资源包共约2000个文件以1998张jpg图像为主体另附1个py脚本与1个json类别文件压缩包大小约39.65MB图像按类别目录分门别类存放便于直接接入ImageFolder等标准加载方式。数据已预先划分为训练集约2200张、验证集约200张、测试集约100张省去自行切分的步骤json文件可用于核对具体类别名称与索引映射。目前已有269人学习下载适合希望快速验证分类网络效果、练习数据增强与迁移学习流程的读者参考使用。1. 垃圾分类图像分类数据集从「这堆图能不能用」到跑通第一个基线你手头可能有一批从社区回收站、校园垃圾桶旁、或者公开渠道攒下来的瓶子与垃圾照片几百到几千张不等文件夹里混着塑料瓶、玻璃瓶、纸盒、果皮命名乱七八糟。现在你想用深度学习做一个图像分类模型把「可回收瓶类」「其他垃圾」这些类别自动分开。这件事的核心不是模型多深而是数据集能不能撑起训练——类别是否均衡、标注是否一致、图像质量是否够用。垃圾分类、瓶子垃圾图像分类这个方向本质是一个细粒度、强背景干扰的视觉分类任务适合刚接触深度学习图像分类的从业者拿来练手也适合做环保、回收分拣场景的工程团队做原型验证。这一章先把「这是什么、能解决什么、适合谁」讲清楚后面几章再落到具体怎么做、参数怎么设、坑在哪。2. 垃圾分类数据集长什么样类别体系、标注格式与选型理由2.1 先定类别体系再谈图像数量很多人一上来就问「多少张图够」但更关键的问题是「分几类」。垃圾分类的类别体系直接决定模型能不能落地。常见做法是两级一级粗分可回收物、厨余垃圾、有害垃圾、其他垃圾二级细分塑料瓶、玻璃瓶、易拉罐、纸盒等。如果你只做瓶子垃圾图像分类建议先聚焦 4 到 6 个视觉差异明显的类比如「透明塑料瓶」「有色塑料瓶」「玻璃瓶」「易拉罐」「纸杯」「其他」。类别太少模型学不到区分度类别太多且视觉边界模糊比如「可回收塑料」和「其他塑料」会让标注一致性崩掉。我一般会先画一张类别定义表把每个类别的判定标准写清楚比如「透明塑料瓶瓶身透光、无标签或标签面积小于 30%」。这张表不是给模型看的是给标注的人看的。没有这张表后面标注出来的数据大概率是玄学。类别名判定要点易混淆项建议最少样本数透明塑料瓶瓶身透光标签面积小玻璃瓶、有色塑料瓶300有色塑料瓶瓶身有颜色不透光透明塑料瓶300玻璃瓶材质反光瓶壁厚透明塑料瓶250易拉罐金属反光圆柱形纸杯250纸杯/纸盒纸质纹理无金属反光易拉罐200其他垃圾不属于以上类别全部200这张表里的「建议最少样本数」不是硬性标准而是经验下限。如果某个类少于 200 张模型大概率会在这个类上翻车。类别不均衡时优先补少样本类而不是继续堆多样本类。2.2 标注格式文件夹分类 vs CSV 清单图像分类最省事的标注方式就是按文件夹分。每个类别一个文件夹文件夹名就是类别名图片直接放进去。这种格式几乎所有训练框架都认比如torchvision.datasets.ImageFolder、tensorflow.keras.preprocessing.image_dataset_from_directory。另一种是 CSV 清单两列filename,label。CSV 的好处是灵活可以随时改标签、做交叉验证划分也方便记录来源和备注。我一般会同时保留两份原始文件夹结构用于快速预览CSV 用于训练脚本读取。下面是一个把文件夹结构转成 CSV 的脚本顺便做一次基础校验。import os import csv from pathlib import Path # 数据集根目录下面每个子文件夹是一个类别 data_root Path(dataset/garbage_cls) # 输出 CSV 路径 csv_path Path(dataset/train_labels.csv) # 允许的图片后缀避免把 .DS_Store 之类的文件读进来 valid_ext {.jpg, .jpeg, .png, .bmp, .webp} rows [] class_names sorted([d.name for d in data_root.iterdir() if d.is_dir()]) print(类别列表:, class_names) for cls_name in class_names: cls_dir data_root / cls_name for img_path in cls_dir.iterdir(): if img_path.suffix.lower() not in valid_ext: continue # 相对路径方便后续迁移数据集位置 rel_path img_path.relative_to(data_root) rows.append({filename: str(rel_path), label: cls_name}) # 写 CSV带表头 with csv_path.open(w, newline, encodingutf-8) as f: writer csv.DictWriter(f, fieldnames[filename, label]) writer.writeheader() writer.writerows(rows) print(f共写入 {len(rows)} 条记录到 {csv_path})这段脚本的逻辑很直白遍历每个类别文件夹过滤掉非图片文件把相对路径和类别名写成 CSV。参数上data_root指向你的数据集根目录csv_path是输出位置。跑完之后你可以用pandas读一下检查每个类别的数量是否均衡。import pandas as pd df pd.read_csv(dataset/train_labels.csv) print(df[label].value_counts())如果某个类别的数量明显偏少比如只有 80 张而其他类有 500 张那就需要先处理不均衡问题而不是直接开训。常见做法是数据增强、重采样、或者用类别权重。这个后面在训练章节会展开。2.3 图像质量与背景干扰垃圾分类数据集的天然难点垃圾分类图像和 ImageNet 那种干净数据集不一样。你的图片可能来自手机拍摄背景里有地面、手、其他垃圾、光照变化。瓶子类还有反光、透明、形变的问题。透明塑料瓶和玻璃瓶在视觉上非常接近模型很容易把两者搞混。这不是模型不行是数据本身就有歧义。我一般会做三件事第一人工抽检 100 张看看有没有明显标错的第二统计图像尺寸分布太小的图比如小于 100x100直接剔除第三对透明瓶和玻璃瓶这类易混类别额外补充侧光、逆光、不同角度的样本。如果条件允许用albumentations做在线增强时对这两类加更强的颜色抖动和随机裁剪让模型学到材质差异而不是背景捷径。提示不要用「瓶子」和「垃圾」这种粗标签直接训。粗标签会让模型学到「有瓶子就是可回收」这种捷径换个场景就废了。3. 用 PyTorch 跑通瓶子垃圾图像分类基线从 DataLoader 到第一个 epoch3.1 环境准备与最小依赖深度学习环境这件事新手最容易卡在版本冲突上。我一般用 conda 建一个干净环境只装必要的东西。下面这套是 CPU 和单卡 GPU 都能跑的配置不追求最新追求能复现。conda create -n garbage_cls python3.10 -y conda activate garbage_cls pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install pandas pillow scikit-learn matplotlib tqdm如果你没有 GPU把cu118换成cpu即可。装完之后用下面这行验证import torch print(torch.__version__, torch.cuda.is_available())torch.cuda.is_available()返回True说明 GPU 可用。返回False也不影响跑基线只是慢一点。3.2 构建 Dataset 和 DataLoaderPyTorch 里最稳的做法是继承torch.utils.data.Dataset自己写__getitem__。这样你能完全控制读取逻辑也方便后面加增强。import torch from torch.utils.data import Dataset, DataLoader from PIL import Image from torchvision import transforms import pandas as pd from pathlib import Path class GarbageDataset(Dataset): def __init__(self, csv_file, img_root, transformNone): self.df pd.read_csv(csv_file) self.img_root Path(img_root) self.transform transform # 类别到索引的映射排序保证可复现 self.classes sorted(self.df[label].unique()) self.class_to_idx {c: i for i, c in enumerate(self.classes)} def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] img_path self.img_root / row[filename] # 统一转 RGB防止灰度图或 RGBA 图导致通道数不一致 image Image.open(img_path).convert(RGB) label self.class_to_idx[row[label]] if self.transform: image self.transform(image) return image, label # 训练增强随机裁剪、翻转、颜色抖动 train_tf transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.3), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) # 验证集只做 resize 和归一化不做随机增强 val_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]), ]) train_ds GarbageDataset(dataset/train_labels.csv, dataset/garbage_cls, train_tf) val_ds GarbageDataset(dataset/val_labels.csv, dataset/garbage_cls, val_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4)这里有几个参数值得说清楚。batch_size32是单卡 8G 显存下的稳妥值显存不够就降到 16 或 8。num_workers4在 Windows 上有时会出问题如果报错就改成 0。Normalize用的均值和方差是 ImageNet 的统计值因为后面要用预训练模型保持一致最省事。RandomResizedCrop的scale(0.7, 1.0)表示随机裁取原图 70% 到 100% 的区域对瓶子这种主体居中的图比较合适裁太狠会把瓶身切掉。3.3 选一个 backboneResNet18 还是更小的模型垃圾分类数据集通常不大几千到几万张。这种规模下ResNet18 或 EfficientNet-B0 是性价比最高的选择。ResNet18 参数量约 1100 万ImageNet 预训练权重容易拿训练稳定。EfficientNet-B0 更小但输入分辨率要求 224 以上对透明瓶的细节可能更友好。我一般先用 ResNet18 跑基线如果准确率卡在某个值上不去再换 backbone 对比。import torch.nn as nn from torchvision import models def build_model(num_classes, backboneresnet18, pretrainedTrue): if backbone resnet18: model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT if pretrained else None) # 替换最后的全连接层输出类别数 model.fc nn.Linear(model.fc.in_features, num_classes) elif backbone efficientnet_b0: model models.efficientnet_b0(weightsmodels.EfficientNet_B0_Weights.DEFAULT if pretrained else None) model.classifier[1] nn.Linear(model.classifier[1].in_features, num_classes) else: raise ValueError(f不支持的 backbone: {backbone}) return model num_classes len(train_ds.classes) model build_model(num_classes, backboneresnet18, pretrainedTrue) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) print(f类别数: {num_classes}, 设备: {device})pretrainedTrue会加载 ImageNet 权重这对小数据集非常关键。从头训 ResNet18 在几千张图上大概率过拟合。替换fc层之后只有最后一层是随机初始化的前面都是预训练特征。3.4 训练循环与关键参数训练循环本身不复杂但有几个参数决定你能不能跑出结果。优化器用AdamW学习率1e-3对预训练模型微调偏大我一般用3e-4或1e-4。损失函数用CrossEntropyLoss如果类别不均衡加weight参数。import torch.optim as optim from tqdm import tqdm criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max20) def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total 0.0, 0, 0 for images, labels in tqdm(loader, desctrain): images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() total_loss, correct, total 0.0, 0, 0 for images, labels in tqdm(loader, descval): images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total best_acc 0.0 for epoch in range(20): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fepoch {epoch1}: train_loss{train_loss:.4f} train_acc{train_acc:.4f} fval_loss{val_loss:.4f} val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_garbage_cls.pth) print(f保存最佳模型val_acc{best_acc:.4f})CosineAnnealingLR的T_max20表示 20 个 epoch 内学习率按余弦曲线降到接近 0。如果你只跑 10 个 epoch就把T_max改成 10。weight_decay1e-4是 AdamW 的常规值能压一点过拟合。保存模型时只存state_dict不存整个模型对象这样加载时更灵活。跑完 20 个 epoch如果验证集准确率能到 85% 以上说明数据集和基线是通的。如果卡在 60% 左右大概率是类别混淆或标注问题不是模型问题。4. 垃圾分类图像分类的避坑与排查标注、不均衡、过拟合、透明瓶混淆4.1 标注不一致同一个瓶子被标成两个类现象训练 loss 震荡验证准确率上不去混淆矩阵里某两个类互相错得特别多。原因标注时没有统一标准比如有人把「透明塑料瓶」标成「玻璃瓶」有人反过来。解决先抽 200 张做一致性检查让两个人独立标同一批图算一下一致率。低于 90% 就重新定标准把易混类别的判定规则写得更细比如「瓶底有厚玻璃纹路算玻璃瓶瓶身有塑料接缝算塑料瓶」。4.2 类别不均衡少样本类被模型忽略现象总体准确率看着还行但少样本类的召回率极低模型几乎全预测成多样本类。原因CrossEntropyLoss默认每个样本权重一样多样本类主导梯度。解决给CrossEntropyLoss加weight权重和类别频率成反比。或者用重采样让每个 batch 里少样本类被抽到的概率更高。import numpy as np class_counts train_ds.df[label].value_counts().sort_index().values class_weights 1.0 / class_counts class_weights class_weights / class_weights.sum() * len(class_counts) class_weights torch.tensor(class_weights, dtypetorch.float32).to(device) criterion nn.CrossEntropyLoss(weightclass_weights) print(类别权重:, class_weights)这段代码先统计每个类的样本数取倒数再归一化到均值为 1 左右。weight传给损失函数后少样本类的 loss 会被放大梯度更新时更受重视。4.3 过拟合训练准确率 99%验证准确率 70%现象训练集准确率一路涨到 99%验证集卡在 70% 不动甚至下降。原因数据集太小、模型太大、增强不够。解决先加数据增强RandomResizedCrop、ColorJitter、RandomRotation都可以上。再加DropoutResNet18 的fc前面可以插一个nn.Dropout(0.5)。最后考虑冻结前面层只训fc或者用更小的学习率。4.4 透明瓶与玻璃瓶混淆模型学的是背景不是材质现象混淆矩阵里透明塑料瓶和玻璃瓶互相错得最多但人眼看这两类差异明显。原因模型可能学到了背景捷径比如透明瓶常出现在某类背景上玻璃瓶在另一类背景上。解决把背景裁掉或者用分割模型先抠出瓶身再分类。更简单的做法是加随机背景替换增强把瓶子贴到不同背景上逼模型看瓶身。4.5 图像尺寸和通道不一致导致训练中断现象训练到一半报错提示 tensor 尺寸不匹配或通道数不对。原因数据集里混了灰度图、RGBA 图、尺寸差异极大的图。解决在__getitem__里统一convert(RGB)在transform里统一Resize。如果还有问题写个脚本扫一遍所有图片把打不开的、尺寸小于 50x50 的直接列出来剔除。from PIL import Image from pathlib import Path bad_files [] for img_path in Path(dataset/garbage_cls).rglob(*): if img_path.suffix.lower() not in {.jpg, .jpeg, .png, .bmp, .webp}: continue try: with Image.open(img_path) as im: w, h im.size if w 50 or h 50: bad_files.append((str(img_path), too_small, w, h)) except Exception as e: bad_files.append((str(img_path), unreadable, str(e))) for item in bad_files: print(item) print(f共发现 {len(bad_files)} 个问题文件)这个脚本会列出所有打不开或太小的图片。跑一遍把问题文件处理掉能省掉很多训练中途报错的麻烦。5. 把基线推到可用混淆矩阵诊断、难例挖掘与一个提点技巧5.1 用混淆矩阵定位问题类别准确率只是一个数字真正有用的是混淆矩阵。下面这段代码在验证集上跑一遍画出每个类别的错分情况。from sklearn.metrics import confusion_matrix, classification_report import numpy as np torch.no_grad() def get_all_preds(model, loader, device): model.eval() all_preds, all_labels [], [] for images, labels in loader: images images.to(device) outputs model(images) preds outputs.argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) return np.array(all_labels), np.array(all_preds) labels, preds get_all_preds(model, val_loader, device) cm confusion_matrix(labels, preds) print(混淆矩阵:) print(cm) print(classification_report(labels, preds, target_namesval_ds.classes))classification_report会给出每个类的 precision、recall、f1-score。如果某个类的 recall 特别低说明这个类被大量错分。结合混淆矩阵你能看出它主要被错分成了哪个类。比如透明塑料瓶的 recall 只有 0.6且大部分错分到玻璃瓶那就针对这两类做难例挖掘。5.2 难例挖掘把错分样本挑出来单独看难例挖掘的思路很简单用当前模型跑一遍训练集把预测错误的样本挑出来人工看一遍确认是标注问题还是模型问题。如果是标注问题修正标签如果是模型问题把这些难例复制几份加入训练集或者对这些类做更强的增强。torch.no_grad() def find_hard_examples(model, dataset, device, top_k50): model.eval() hard [] for idx in range(len(dataset)): image, label dataset[idx] # dataset 返回的是增强后的图这里为了诊断用原始图更合适 # 简化处理直接用增强后的图推理 input_tensor image.unsqueeze(0).to(device) output model(input_tensor) pred output.argmax(dim1).item() if pred ! label: hard.append((idx, label, pred)) return hard[:top_k] hard_examples find_hard_examples(model, val_ds, device, top_k30) for idx, true_label, pred_label in hard_examples: print(fidx{idx} 真实{val_ds.classes[true_label]} 预测{val_ds.classes[pred_label]})这段代码会列出验证集里被错分的样本索引和类别。你可以根据索引找到原图看看是标注错了还是模型确实分不出来。如果是标注错了改标签如果是模型问题把这些图加入训练集或者对这类图做针对性增强。5.3 一个提点技巧两阶段训练与标签平滑如果基线准确率卡在 85% 左右上不去可以试两阶段训练。第一阶段冻结 backbone只训分类头学习率设1e-3跑 5 个 epoch。第二阶段解冻全部层学习率降到1e-4跑 15 个 epoch。这样能让分类头先适应数据分布再微调整个网络比一上来就全量微调更稳。另外标签平滑label smoothing对垃圾分类这种有歧义的类别很有效。CrossEntropyLoss(label_smoothing0.1)会让模型不要把概率压到 1.0对易混类别更宽容。我一般会在第二阶段加上通常能提 1 到 2 个点。# 第二阶段解冻全部层小学习率微调 for param in model.parameters(): param.requires_grad True optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) criterion nn.CrossEntropyLoss(label_smoothing0.1) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max15) for epoch in range(15): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fstage2 epoch {epoch1}: val_acc{val_acc:.4f})标签平滑的0.1是经验值太大比如 0.3会让模型欠拟合太小比如 0.01效果不明显。两阶段训练的关键是第二阶段学习率一定要小否则会破坏预训练特征。我自己的习惯是每次跑完基线先看混淆矩阵再决定是补数据、改标注还是调模型。数据集的问题永远比模型多。希望帮到你。本文还有配套的精品资源点击获取