ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

8万张图245个类:垃圾分类数据集与TF代码实战指南

8万张图245个类:垃圾分类数据集与TF代码实战指南 简介这份资源面向深度学习入门者与计算机视觉方向的开发者提供一套完整的垃圾分类图像分类实战方案包含约8万张图片、覆盖245个类别的数据集以及基于TensorFlow实现的训练与推理代码可用于课程设计、竞赛练手或迁移学习实验。压缩包共73个文件约101.99MB其中jpg、png、jpeg等图片文件构成数据集主体py脚本负责数据读取、CNN与MobileNet训练及模型测试h5与pickle保存训练好的权重另有txt、xml、md等说明与标注文件models目录下附带两组已训练模型可直接加载验证效果。目前已有29958人学习下载热度较高。读者可据此复现从数据组织、模型搭建到训练评估的完整流程并借助现成权重快速开展推理与二次调参目录结构清晰便于按模块查阅与替换数据集。1. 8万张图、245个类垃圾分类数据集和TF代码到底能跑出什么手里有一份「垃圾分类数据集和tf代码-8w张图片245个类.zip」第一反应通常不是兴奋而是犯嘀咕8万张图、245个类这规模放在分类任务里不算小但类别数这么细到底是按什么粒度分的是「可回收/厨余/有害/其他」四分类的细分版还是直接按具体物品名拆成了两百多类这个判断会直接决定你后面用不用得上、要不要重新合并标签。这份资源的价值在于它同时给了数据和 TensorFlow 训练代码省掉了从零搭 pipeline 的时间。适合两类人一类是想快速验证垃圾分类模型能不能落地到边缘设备的工程师另一类是拿它当多分类练手项目、想搞清楚大数据集怎么喂进 TF 的学生。但 245 个类意味着类别极不均衡长尾类可能只有几十张图直接开训大概率翻车。下面按「先看清数据 → 再跑通代码 → 再调参 → 再避坑」的顺序讲透。2. 先摸清数据底细245个类怎么分布、图片什么规格拿到压缩包别急着解压完就开训先花二十分钟把数据摸清楚。这一步偷懒后面训练曲线会教你做人。2.1 解压后的目录结构判断常见做法是解压后看到类似dataset/train/类别名/xxx.jpg的结构也可能是images/加一个labels.csv的扁平结构。两种处理方式完全不同先确认属于哪种。# 先看压缩包内顶层结构不解压 unzip -l 垃圾分类数据集和tf代码-8w张图片245个类.zip | head -50 # 解压到指定目录 unzip -q 垃圾分类数据集和tf代码-8w张图片245个类.zip -d garbage_data # 看目录层级 find garbage_data -maxdepth 3 -type d | head -30unzip -l先列清单是为了避免解压出一堆嵌套目录后找不到北。-q静默解压-d指定目标目录。如果find出来的目录层级超过三层说明里面可能还套了一层压缩包或者按日期分了子目录需要再处理。2.2 统计每个类的图片数量和尺寸分布245 个类必须用脚本统计肉眼看不完。下面这段统计类别分布和图片尺寸直接决定后面要不要做重采样和统一 resize。import os from collections import Counter from PIL import Image root garbage_data/dataset/train class_counts {} size_counter Counter() for cls in os.listdir(root): cls_dir os.path.join(root, cls) if not os.path.isdir(cls_dir): continue imgs [f for f in os.listdir(cls_dir) if f.lower().endswith((.jpg, .jpeg, .png))] class_counts[cls] len(imgs) # 抽样看尺寸不用全量读 for f in imgs[:5]: try: with Image.open(os.path.join(cls_dir, f)) as im: size_counter[im.size] 1 except Exception: pass counts sorted(class_counts.values()) print(类别总数:, len(class_counts)) print(最少类样本数:, counts[0]) print(最多类样本数:, counts[-1]) print(中位数:, counts[len(counts)//2]) print(尺寸Top5:, size_counter.most_common(5))这段代码的关键输出是三个数最少类样本数、最多类样本数、中位数。如果最少类只有个位数、最多类上千那类别不均衡非常严重后面 loss 和采样策略必须动。尺寸统计用抽样而非全量是因为 8 万张图全读一遍太慢抽前 5 张足够看出规格是否统一。2.3 判断要不要合并类别245 个类如果长尾严重一个务实的选择是把样本数低于某个阈值的类合并成「其他」类或者按材质大类归并。常见做法是设一个阈值比如少于 100 张的类先看它们是不是语义相近相近就合并。判断维度处理建议最少类 50 张合并到相近大类或直接剔除类别间样本比 50:1必须用加权 loss 或重采样图片尺寸不统一训练前统一 resize别指望 TF 自动处理存在灰度图混彩色图统一转 RGB否则通道数对不上提示合并类别前先备份原始标签映射否则后面想还原细分粒度就麻烦了。3. 把TF代码跑起来从数据管道到第一个baseline数据摸清了接下来让 TensorFlow 代码跑通。这一步的目标不是刷高准确率而是先拿到一个能跑完、能出数的 baseline确认整条链路没断。3.1 用 tf.data 构建高效输入管道8 万张图如果不用tf.data而用ImageDataGeneratorGPU 会大量时间空等。下面用image_dataset_from_directory快速搭管道这是 TF 里最省事的入口。import tensorflow as tf IMG_SIZE (224, 224) BATCH_SIZE 64 AUTOTUNE tf.data.AUTOTUNE train_ds tf.keras.utils.image_dataset_from_directory( garbage_data/dataset/train, image_sizeIMG_SIZE, batch_sizeBATCH_SIZE, label_modeint, shuffleTrue, seed42, ) val_ds tf.keras.utils.image_dataset_from_directory( garbage_data/dataset/val, image_sizeIMG_SIZE, batch_sizeBATCH_SIZE, label_modeint, shuffleFalse, ) # 预取和缓存避免IO成为瓶颈 train_ds train_ds.cache().prefetch(buffer_sizeAUTOTUNE) val_ds val_ds.cache().prefetch(buffer_sizeAUTOTUNE) # 拿到类别数用于输出层 num_classes len(train_ds.class_names) print(类别数:, num_classes)image_dataset_from_directory会自动按子目录名生成标签label_modeint返回整数标签配合sparse_categorical_crossentropy。cache()把数据缓存在内存或本地8 万张图如果内存放不下会落盘prefetch让 CPU 预处理和 GPU 计算重叠。num_classes必须从数据里读不能写死 245因为合并类别后这个数会变。3.2 搭一个能跑通的分类网络baseline 不用上太复杂的结构先用迁移学习拿 MobileNetV2 或 EfficientNet 试水收敛快、显存友好。from tensorflow.keras import layers, models base tf.keras.applications.MobileNetV2( input_shape(224, 224, 3), include_topFalse, weightsimagenet, ) base.trainable False # 先冻结只训分类头 model models.Sequential([ layers.Rescaling(1./127.5, offset-1), # MobileNetV2 要求 [-1,1] base, layers.GlobalAveragePooling2D(), layers.Dropout(0.3), layers.Dense(num_classes, activationsoftmax), ]) model.compile( optimizertf.keras.optimizers.Adam(1e-3), losssparse_categorical_crossentropy, metrics[accuracy], ) model.summary()include_topFalse去掉原分类头trainableFalse先冻结主干只训练后面的 Dense 层。Rescaling把 [0,255] 映射到 [-1,1]这是 MobileNetV2 预训练时的输入要求漏了这步精度会明显掉。Dropout(0.3)在 245 类大数据集上帮助有限但没坏处。输出层节点数用num_classes激活用 softmax 配稀疏交叉熵。3.3 处理类别不均衡的损失函数245 个类不均衡是常态直接训会让模型偏向头部类。用类别权重是最省事的补救。import numpy as np # 统计每个类的样本数 class_counts np.bincount( [y for _, y in train_ds.unbatch().map(lambda x, y: (x, y))] ) total class_counts.sum() # 权重与频次成反比做归一化 class_weight {i: total / (len(class_counts) * c) for i, c in enumerate(class_counts) if c 0} history model.fit( train_ds, validation_dataval_ds, epochs10, class_weightclass_weight, )np.bincount统计每个类出现次数class_weight让稀有类的样本在 loss 里占更大权重。注意class_weight的 key 必须是整数标签且要跳过计数为 0 的类否则会除零。如果某些类样本极少权重会非常大可以设一个上限比如 50避免训练被少数类带偏。注意train_ds.unbatch()遍历一遍 8 万张图比较慢可以在第一次统计后把class_weight存成 json后续直接读。4. 参数怎么调让245类分类从能跑到能用baseline 跑通后准确率大概率不好看。这一章讲几个真正影响结果的参数和策略都是踩过坑总结出来的。4.1 输入分辨率与batch size的取舍224×224 是默认值但垃圾分类里很多判别靠纹理和细节比如塑料瓶和玻璃瓶的区分。分辨率提到 320 或 384 通常能涨点代价是显存和训练时间。分辨率显存占用相对适用场景224×2241x快速验证、边缘部署320×320约2x精度优先、显存够384×384约3x细粒度分类、离线推理batch size 方面64 是稳妥起点。显存不够就降到 32但要注意 batch 变小后学习率也要相应下调否则训练不稳定。常见做法是 batch 翻倍时学习率也翻倍反过来同理。4.2 解冻主干与分层学习率冻结主干训几个 epoch 后分类头基本收敛这时候解冻部分主干层做微调精度还能再上一截。关键是学习率要小。# 解冻最后30层做微调 base.trainable True for layer in base.layers[:-30]: layer.trainable False model.compile( optimizertf.keras.optimizers.Adam(1e-5), # 微调用小学习率 losssparse_categorical_crossentropy, metrics[accuracy], ) fine_tune_history model.fit( train_ds, validation_dataval_ds, epochs10, class_weightclass_weight, )解冻后学习率从 1e-3 降到 1e-5是因为预训练权重已经很好大学习率会把它们破坏掉。只解冻最后 30 层而不是全部是为了平衡效果和过拟合风险。如果验证集准确率在解冻后反而下降说明解冻太多或学习率还是偏大。4.3 数据增强的度怎么把握垃圾分类的图片通常是物品正面照旋转和翻转是安全的增强但颜色抖动要谨慎因为颜色本身是分类依据比如绿色厨余、蓝色可回收。data_aug tf.keras.Sequential([ layers.RandomFlip(horizontal), layers.RandomRotation(0.1), layers.RandomZoom(0.1), # 颜色抖动幅度要小避免破坏颜色语义 layers.RandomBrightness(0.1), layers.RandomContrast(0.1), ]) train_ds_aug train_ds.map( lambda x, y: (data_aug(x, trainingTrue), y), num_parallel_callsAUTOTUNE, ).prefetch(AUTOTUNE)RandomFlip和RandomRotation幅度小不会让物品变得不可识别。RandomBrightness和RandomContrast控制在 0.1是因为垃圾分类里颜色是强特征抖太狠等于给模型制造矛盾样本。增强层放在map里而不是模型里是为了验证集不受影响。提示增强后一定要肉眼抽查几张图确认没有把标签语义破坏掉这一步很多人跳过结果训出来模型对颜色不敏感。5. 避坑与排查8万张图训练时最容易翻车的5个地方这一章全是血泪经验每条都按「现象 → 原因 → 解决」写遇到对应症状直接对号入座。5.1 训练loss不降准确率卡在1%左右现象第一个 epoch 结束 loss 几乎不动准确率约等于 1/245。原因最常见是标签没对上。image_dataset_from_directory按目录名排序生成整数标签如果你自己另外维护了一份标签映射顺序很可能不一致。另一个原因是输入没做Rescaling像素值 [0,255] 直接喂进预训练模型。解决打印train_ds.class_names确认标签顺序和你的映射表逐一对齐。检查模型第一层是不是RescalingMobileNetV2 必须把输入映射到 [-1,1]。5.2 验证集准确率远高于训练集现象训练集准确率 60%验证集 85%反常。原因验证集太小或者和训练集有重叠。8 万张图如果验证集只切了几百张且恰好都是简单类就会出现这种假象。另一个可能是验证集没做 shuffle 但分布恰好集中。解决验证集至少占 10%且按类别分层抽样。用train_test_split的stratify参数保证每个类在验证集里都有代表。5.3 训练到一半loss突然变NaN现象前几个 epoch 正常突然 loss 变成 nan权重全废。原因学习率太大或者class_weight里某些类权重过高导致梯度爆炸。245 类里如果有类只有几张图权重可能上百。解决给class_weight设上限比如min(weight, 50)。同时加梯度裁剪optimizertf.keras.optimizers.Adam(1e-3, clipnorm1.0)。如果已经 NaN只能从头训没有后悔药。5.4 GPU利用率低训练速度慢现象nvidia-smi显示 GPU 利用率在 20% 以下波动。原因数据管道是瓶颈CPU 预处理跟不上 GPU。8 万张图如果每次 epoch 都重新解码IO 压力很大。解决cache()到本地 SSDprefetch(AUTOTUNE)必须加num_parallel_calls设成 CPU 核心数。如果内存够大cache()不加参数直接缓内存最快。5.5 保存的模型加载后预测结果全乱现象训练时准确率 90%保存成.h5再加载预测全是同一类。原因保存时没保存预处理层或者加载后输入格式和训练时不一致。Rescaling层如果在模型外保存的模型就不含它。解决把Rescaling放进Sequential模型内部这样保存时一起存。加载后用和训练时完全相同的预处理流程包括 resize 的插值方法。注意排查问题时先固定随机种子否则每次结果不一样根本没法定位是改动生效还是随机波动。6. 进阶技巧用混淆矩阵和Top-5准确率看清245类的真实表现245 类分类单看 overall accuracy 会骗人。头部几十个类可能贡献了大部分正确样本长尾类全军覆没但被平均值掩盖。这一章讲两个能看清真相的方法以及一个我常用的验证习惯。6.1 混淆矩阵定位「重灾区」不要画 245×245 的全矩阵看不清。挑出错误最多的前 20 个类看它们被错分到哪去了。import numpy as np import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix import seaborn as sns # 收集验证集预测 y_true, y_pred [], [] for x, y in val_ds: preds model.predict(x, verbose0) y_true.extend(y.numpy()) y_pred.extend(np.argmax(preds, axis1)) y_true np.array(y_true) y_pred np.array(y_pred) # 找出错误最多的类 errors_per_class {} for c in range(num_classes): mask y_true c if mask.sum() 0: continue errors_per_class[c] (y_pred[mask] ! c).sum() top_error_classes sorted(errors_per_class, keyerrors_per_class.get, reverseTrue)[:20] cm confusion_matrix(y_true, y_pred, labelstop_error_classes) plt.figure(figsize(14, 12)) sns.heatmap(cm, annotFalse, cmapBlues, xticklabelstop_error_classes, yticklabelstop_error_classes) plt.xlabel(Predicted) plt.ylabel(True) plt.tight_layout() plt.savefig(confusion_top20.png, dpi150)这段代码先跑一遍验证集收集预测再统计每个类的错误数取前 20 个画热力图。对角线亮说明该类本身还行非对角线亮说明两类容易混。垃圾分类里常见的混淆对是「纸杯 vs 塑料杯」「玻璃瓶 vs 塑料瓶」看矩阵能直接定位。6.2 Top-5准确率作为辅助指标245 类里模型把「塑料瓶」排第一错成「玻璃瓶」但排第二是对的这种错误在实际业务里可以接受。Top-5 准确率能反映模型的「大致方向」对不对。top5 tf.keras.metrics.SparseTopKCategoricalAccuracy(k5) for x, y in val_ds: preds model.predict(x, verbose0) top5.update_state(y, preds) print(Top-5 准确率:, top5.result().numpy())SparseTopKCategoricalAccuracy直接吃整数标签和 softmax 输出不用 one-hot。如果 Top-1 是 70% 而 Top-5 是 95%说明模型大部分时候能框对方向只是细粒度判别不够这时候可以考虑用更细的注意力机制或者更高分辨率而不是盲目加数据。6.3 我自己的验证习惯每次改完参数我不会只看一个数字。固定验证集、固定随机种子跑完存三样东西混淆矩阵图、Top-5 准确率、每个类的召回率列表。然后和上一次的结果对比看是整体涨了还是只有头部类涨了。如果只有头部类涨、长尾类没动说明改动没解决根本问题。还有一个习惯是训练前先拿 100 张图过一遍模型看预测分布是不是均匀的。如果 100 张图全预测成同一类说明初始化或标签有问题不用等训完 10 个 epoch 才发现。这套流程走下来8 万张图 245 类的垃圾分类任务基本能从「跑通」走到「能用」。数据集和 TF 代码只是起点真正决定效果的是你对数据分布的理解和参数调整的耐心。希望帮到你。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进