8万张图245类垃圾分类数据集：TensorFlow图像分类完整实战

简介面向垃圾分类图像识别任务这套资料包含八万张图片、二百四十五个类别的数据集同时提供可直接运行的深度学习训练代码下载后即可用于模型训练与评估适合深度学习初学者、参加算法竞赛的学生以及需要快速搭建垃圾分类模型的开发者。资源包内共两千个文件整包约561.25MB。其中以JPG和JPEG图片为主覆盖电池、塑料瓶、包装盒、玻璃瓶等常见废弃物类别另有十三个脚本文件、一个类别配置文件和一个说明文档分别对应模型训练、类别映射与使用说明。目前已有一百三十一人浏览学习。对缺少现成数据集的毕业设计或实验项目而言这套资料省去了自行采集、清洗和标注图片的环节解压后即可按文档步骤展开训练与调参是一份能直接上手的完整实践素材。1. 8万张图片、245个类别这套垃圾分类数据包把“最难的部分”提前解决了做垃圾分类识别最头疼的不是模型结构而是数据。社区垃圾亭想装智能识别模块工厂分拣线想按材质自动分流问到最多的就是网上开源数据集不是几千张的玩具就是类目太粗分不清矿泉水瓶和饮料瓶。这套 8w 张图片、245 个类别的垃圾分类数据集连同 tf 代码一起打包解压就能开始训练正好补上“能直接用”这个缺口。对做智能分拣设备、垃圾巡检机器人或者拿图像分类练手的工程师和学生来说它能省掉至少一两个月的收集清洗时间。8w 张是带类别标注的原始图像数据不是那种只有几十张的演示 demo245 类覆盖日常生活常见可回收物、厨余、有害垃圾和其他垃圾的细粒度物品后面接硬件分拣时这个粒度的信息才够用。2. 数据集的真实构成目录结构、图片规格与标签核对方法2.1 一份“下载即用”的数据包解压后通常长这样拿到数据包第一件事不是写训练脚本而是把目录结构摸清楚。我见过太多人上来就写model.fit结果跑了半天发现类别数对不上或者 labels.txt 顺序和子目录名不一致白白浪费一下午。正确的顺序是先解压、再统计、最后才碰代码。unzip waste_dataset.zip -d waste_dataset cd waste_dataset tree -L 2 -d常见的打包方式是train/、val/有的包叫test/按类别建子目录同时附带一个labels.txt或class_names.txt记录类别顺序。解压后的典型文件布局如下表这套结构好处是能直接被tf.keras.utils.image_dataset_from_directory读取不需要自己写解析器。路径内容说明train/按 245 个类名建子目录每类若干张 jpg训练集直接喂给目录读取器val/与 train 相同的结构类目一一对应验证集用于早停和挑选模型test/部分数据包会有留作最终评估不要让它参与训练labels.txt每行一个类名顺序必须和子目录名一致推理时按行号映射图片文件本身规格不统一是这类真实数据集的特点。常见的是 JPEG 格式分辨率从几百像素到两三千像素都有。8w 张图如果全是原始分辨率总占用会到 GB 量级所以看到压缩包十几 GB 是正常的不用怀疑下载坏了。真正需要警惕的反而是包内文件数对不上、某些类目录为空这类问题。2.2 245 类怎么组织从四分类到细粒度物品垃圾分类的经典粗分法是可回收物、厨余、有害垃圾、其他垃圾四类。245 类相当于把这四类继续往下拆可回收物里有纸板箱、矿泉水瓶、易拉罐、玻璃瓶、旧衣物厨余里有菜叶、苹果核、剩饭有害垃圾里有废电池、过期药品、灯管。一个常见的设计是直接在类名里体现层级比如recyclable_plastic_bottle、kitchen_waste_apple_core这种命名法。这样做的好处是训练时直接当 245 类多分类处理但评估时可以按类名前缀随时聚合成四类接到后端分拣设备时也能拿到细粒度信息——你想让机械臂把矿泉水瓶和饮料瓶分别抓进不同料仓就需要这个粒度。类目拆得细也带来问题很多类在视觉上高度相似比如不同品牌的矿泉水瓶、不同颜色的易拉罐人眼都要仔细分辨。所以这类数据集的标注质量比类目数量更关键下载后先抽样检查比急着训练重要得多。2.3 训练前先做标签核对下载即用的数据不等于可以直接信任。我一般会先跑两个脚本第一个统计文件数量和每类样本分布第二个检查图片能否正常解码。import os root waste_dataset/train class_names sorted(os.listdir(root)) counts { c: len(os.listdir(os.path.join(root, c))) for c in class_names } print(类别数:, len(class_names)) print(图片总数:, sum(counts.values())) print(最少样本类别:, min(counts, keycounts.get)) print(最少样本数:, counts[min(counts, keycounts.get)]) print(最多样本类别:, max(counts, keycounts.get)) print(最多样本数:, counts[max(counts, keycounts.get)])如果统计出来总数在 8w 附近、245 个类都有样本说明文件层面的完整性没问题。接着用 OpenCV 或 TensorFlow 随机解码几十张图把解码失败的路径打印出来同时把labels.txt的每一行和train/下的目录名做差集比对。类目顺序和目录名不一致的情况在网上下载的数据包里经常出现训练脚本里一旦按行号读标签这个不一致会让验证集准确率看起来永远上不去而这个坑不在模型在标签映射。3. 把 tf 代码跑起来最小复现命令与 train.py 关键参数3.1 环境准备TensorFlow 2.x 依赖与显存预检这套代码基于 TensorFlow 2.x 是当前最稳妥的假设因为 TF 1.x 的 API 在 2.0 之后改动太大新包不会再用tf.Session那套写法。先确认环境版本再做一次显存预检避免训练中途发现 CUDA 驱动不匹配。python -c import tensorflow as tf; print(tf.__version__) python -c from tensorflow.python.client import device_lib; print(device_lib.list_local_devices())第一行输出 2.x 版本号就说明环境可用第二行会列出 GPU。如果是 NVIDIA 显卡能看到name: /device:GPU:0类似输出。没有 GPU 也不是不能训但 8w 张图、245 个类用纯 CPU 训练一轮就得几小时建议至少有一块 8GB 显存的显卡。显存不足的话后面第 5 章有专门的排查思路。提示先跑这两条命令再解压数据否则环境问题会被误判成数据问题排查方向就偏了。3.2 最小训练命令从解压到跑出第一个模型假设数据包已经按第 2 章的布局放好训练入口一般是一个train.py里面用argparse接收数据路径和超参数。最小复现命令通常长这样python train.py \ --data_dir waste_dataset \ --batch_size 32 \ --epochs 50 \ --image_size 224 \ --backbone efficientnet-b0 \ --lr 1e-3先解释一下为什么这么选参数。image_size定为 224 是因为 EfficientNet-B0 和 ResNet50 的预训练权重都要求这个输入尺寸直接用 ImageNet 权重做迁移学习是最省事的路子。batch_size 32是 8GB 显存下的安全值换成 16GB 显存可以提到 64。lr 1e-3是迁移学习的常规起点配合余弦退火或 ReduceLROnPlateau 都能收敛。3.3 train.py 骨架迁移学习、学习率与断点续训一个能直接跑的train.py骨架长这样我在关键位置写了注释。import argparse import tensorflow as tf def build_model(num_classes, backbone_nameefficientnet-b0): base tf.keras.applications.EfficientNetB0( include_topFalse, weightsimagenet, input_shape(224, 224, 3), poolingavg ) base.trainable False # 先冻结只训练分类头 model tf.keras.Sequential([ base, tf.keras.layers.Dropout(0.3), tf.keras.layers.Dense(num_classes, activationsoftmax) ]) return model def main(): parser argparse.ArgumentParser() parser.add_argument(--data_dir, typestr, requiredTrue) parser.add_argument(--batch_size, typeint, default32) parser.add_argument(--epochs, typeint, default50) parser.add_argument(--image_size, typeint, default224) parser.add_argument(--lr, typefloat, default1e-3) args parser.parse_args() train_ds tf.keras.utils.image_dataset_from_directory( args.data_dir /train, image_size(args.image_size, args.image_size), batch_sizeargs.batch_size, label_modeint ) val_ds tf.keras.utils.image_dataset_from_directory( args.data_dir /val, image_size(args.image_size, args.image_size), batch_sizeargs.batch_size, label_modeint ) model build_model(245) model.compile( optimizertf.keras.optimizers.Adam(learning_rateargs.lr), losstf.keras.losses.SparseCategoricalCrossentropy(), metrics[accuracy] ) callbacks [ tf.keras.callbacks.ReduceLROnPlateau(patience3, factor0.5), tf.keras.callbacks.EarlyStopping(patience8, restore_best_weightsTrue), tf.keras.callbacks.ModelCheckpoint( best_model.keras, save_best_onlyTrue ) ] model.fit(train_ds, validation_dataval_ds, epochsargs.epochs, callbackscallbacks) model.save(final_model.keras) if __name__ __main__: main()base.trainable False这行是迁移学习的核心先用 ImageNet 预训练权重把底层特征固定住只训练最后的全连接层。等这一轮训练收敛后再解冻部分底层用更小学习率微调这是处理 245 类小样本场景最稳的路线。label_modeint对应整数标签配合SparseCategoricalCrossentropy用这样不需要做 one-hot 编码省显存也省内存。ReduceLROnPlateau在验证集指标 3 轮不涨时把学习率砍半EarlyStopping连续 8 轮不涨就停并恢复最优权重ModelCheckpoint保存验证集表现最好的模型。这三个回调是训练这类数据集的标配特别是restore_best_weightsTrue不设置的话训练结束拿到的权重不一定是最优的那一份。4. 数据加载管线用 tf.data 喂饱 8w 张图片4.1 为什么不能一次性读进内存流式读取与 cache 的位置8w 张图片、每张几百 KB 到几 MB全量载入内存直接能把 32GB 内存吃光。所以数据加载管线要用流式方案每轮训练只从磁盘读当前 batch 的图片而不是先把全部图片解码成数组。TensorFlow 2.x 的tf.data是官方推荐的组件配合image_dataset_from_directory可以省掉自己写 generator 的麻烦。但流式读取有一个问题每个 epoch 都要重新解码一遍图片IO 开销很大。解决办法是cache()。这里有个关键细节第一次遍历时cache()会把解码后的张量缓存在内存或磁盘上之后每个 epoch 直接从缓存拿数据训练速度能快好几倍。8w 张图解码后的张量大概占十几个 GB如果内存紧张传给cache(dataset_cache)会把缓存写到本地文件。4.2 从目录到批次resize、归一化与 prefetch光靠image_dataset_from_directory还不够工程师一般会在它外面再包一层tf.data管线把归一化、打乱、预取全部串起来。def prepare_dataset(data_dir, batch_size32, image_size224, shuffleTrue): ds tf.keras.utils.image_dataset_from_directory( data_dir, image_size(image_size, image_size), batch_sizebatch_size, label_modeint, shuffleshuffle ) ds ds.map(lambda x, y: (x / 255.0, y)) # 像素归一化到 [0, 1] if shuffle: ds ds.shuffle(1000, reshuffle_each_iterationTrue) ds ds.prefetch(tf.data.AUTOTUNE) return dsmap(lambda x, y: (x / 255.0, y))把像素值从 0-255 缩放到 0-1这是迁移学习的硬性要求——ImageNet 预训练权重就是用这个范围训练的。shuffle(1000)的缓冲区设为 1000 张图保证每个 batch 里类别分布足够随机设太小会让同一类的图片连续出现在一个 batch 里影响梯度稳定性。prefetch(tf.data.AUTOTUNE)让数据加载和 GPU 计算并行GPU 在算当前 batch 时CPU 已经在准备下一个 batch这一步能把 GPU 利用率从 60% 拉到 90% 以上。注意image_dataset_from_directory默认按目录名排序分配标签目录名顺序变化会导致标签顺序变化所以labels.txt必须与目录名排序一致这一点在验证集和测试集上同样成立。4.3 数据增强翻转、裁剪、色彩抖动参数怎么给垃圾图片的拍摄角度、光线、背景差异非常大。同一个矿泉水瓶在阳光下、垃圾桶里、传送带上拍出来差别很大。数据增强是提升泛化能力最便宜的手段TensorFlow 2.x 里可以直接在模型前面加增强层。data_augmentation tf.keras.Sequential([ tf.keras.layers.RandomFlip(horizontal), tf.keras.layers.RandomRotation(0.1), tf.keras.layers.RandomZoom(0.1), tf.keras.layers.RandomContrast(0.1), ])增强参数有三个原则。第一RandomFlip只开水平翻转不要开垂直翻转——垃圾图片倒过来是违背物理常识的。第二旋转角度控制在 10 度以内旋转太多会让小物体标签区域溢出画面。第三对比度抖动幅度 0.1 就够了调太大颜色失真反而破坏迁移学习依赖的纹理特征。增强层放在主干网络之前训练时自动生效推理时不需要做任何改动代码里保留这个设计就行。5. 踩坑与排查从下载到收敛的 5 个真实翻车现场5.1 解压后图片打不开训练中断在 decode 阶段现象是训练跑到一半报InvalidArgumentError: Invalid JPEG data或者某些 epoch 的 loss 突然变成 NaN。原因是压缩包在传输或解压过程中有文件损坏JPEG 文件头缺失或截断。解决分两步先把数据全部解压后做一次完整性巡检把坏图删掉或移到单独的目录。import os from PIL import Image bad_images [] for cls in os.listdir(waste_dataset/train): cls_path os.path.join(waste_dataset/train, cls) for fname in os.listdir(cls_path): fpath os.path.join(cls_path, fname) try: with Image.open(fpath) as img: img.verify() # 只校验文件头不解码全图速度快 except Exception: bad_images.append(fpath) print(f坏图数量: {len(bad_images)}) for p in bad_images[:10]: print(p)img.verify()只解析文件头不把整张图加载进内存跑 8w 张图也就一两分钟。遇到坏图直接从训练目录移走或者用tf.io.decode_image加on-error处理跳过。这一步做完训练中断的问题基本消失。5.2 8w 张图把内存吃光机器直接卡死现象是训练开始后内存占用一路飙到 32GB 甚至更多系统开始疯狂换页训练速度掉到龟速。原因是map里做了归一化cache()把解码后的 8w 张图全部缓存进了内存。解决方法是把缓存改为磁盘缓存或者干脆去掉cache()用原始读取 prefetch的方式。ds ds.cache(waste_dataset_cache) # 缓存到磁盘而不是内存磁盘缓存的代价是首次遍历会写文件多花几分钟但后续每个 epoch 直接从本地缓存读图速度依然远快于重新解码。如果磁盘空间也紧张就把cache()位置放在归一化之前只缓存原始解码结果内存消耗能减半——不过显存和内存都不宽裕的话直接去掉cache()用更合理的prefetch(AUTOTUNE)也能接受只是每个 epoch 多花点解码时间。5.3 显存 OOMCUDNN 报错让人摸不着头脑现象是 NVIDIA 驱动正常但一跑fit就报ResourceExhaustedError: OOM when allocating tensor有时候是Could not create cudnn handle。原因是 batch size 太大、模型输入分辨率太高或别的进程占了显存。解决第一步是设置显存按需增长让 TensorFlow 不用一次性占满整块显卡。physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_growth(physical_devices[0], True)set_memory_growth(True)让显存按需分配避免和别的进程抢资源。这样改完还 OOM就把batch_size从 64 降到 32再从 32 降到 16。训练这种大分类任务batch size 减小带来的准确率损失一般比显存溢出导致的中断要小得多。5.4 245 个类里总有几十个类样本特别少现象是训练 loss 能降但那些样本少的类别 recall 几乎为 0。原因是长尾分布某些类别几百张某些类别只有几十张模型为了把整体 loss 降下来会自动忽略少数类。解决用类别权重让少数类在 loss 里占更大比重。from sklearn.utils.class_weight import compute_class_weight y_labels ... # 从 train_ds 收集全部标签 weights compute_class_weight(balanced, classesnp.unique(y_labels), yy_labels) class_weight dict(enumerate(weights)) model.fit(train_ds, class_weightclass_weight, ...)compute_class_weight(balanced)会让样本数少的类获得更大的权重总样本量除以类别数再除以每类样本数公式很简单但效果很明显。少数类准确率低的另一个补救是数据增强里对少数类多做随机裁剪但 class_weight 是最快见效的改动。5.5 tf 版本差异旧代码在新环境下的兼容处理现象是网上搜到的老代码报AttributeError: module tensorflow has no attribute Session或placeholder这通常是 tf 1.x 时代的脚本。解决分两条路老代码在 tf 2.x 下加import tensorflow.compat.v1 as tf并执行tf.disable_v2_behavior()但这只是临时续命GPU 支持和新 API 都用不顺。更推荐直接把tf.Session()换成model.fit把tf.placeholder换成keras.Input工作量通常只有一两个小时。这套 8w 张的数据包既然标注“下载即用”里面配的代码大概率已经是 2.x 写法但下载后先grep -r tf.Session .查一遍省得训练到一半才发现用的是老接口。6. 验证与落地单图推理、混淆矩阵与模型导出6.1 用训练好的权重跑一遍单图推理训练结束后的第一件事不是看训练集准确率而是拿一张训练集之外的真实照片跑推理。import tensorflow as tf import numpy as np from PIL import Image model tf.keras.models.load_model(best_model.keras) img Image.open(real_garbage.jpg).resize((224, 224)) x np.array(img) / 255.0 x np.expand_dims(x, axis0) pred model.predict(x)[0] top5 np.argsort(pred)[-5:][::-1] for i in top5: print(f类别 {i}: {pred[i]:.3f})推理前先打开labels.txt把行号和这里的i对应起来看前五名里有没有合理的类别。如果一张矿泉水瓶图片预测出来的前五名全是厨余垃圾不要急着调模型先检查labels.txt的顺序是不是和image_dataset_from_directory的目录排序一致——这个坑在验证阶段出现的频率远超想象。6.2 导出 SavedModel给部署留一条后路训练完成后模型还在.keras格式里这种格式适合继续训练不适合部署。对接到 TensorFlow Serving、TFLite 或树莓派这种边端设备时统一导出成 SavedModel 格式更省事。model.export(saved_model_dir)导出后目录里会有saved_model.pb和variables/文件夹推理端可以直接用tf.saved_model.load加载。我的习惯是训练完顺手导出不等到要部署时再找代码重跑一遍——模型保存路径、GPU 环境、训练代码版本一旦换了复现成本会非常高。这个数据集加配套 tf 代码的价值说到底就是把“数据收集、清洗、标注、代码调试”这几件最耗时的杂活提前做完。拿到手后按第 2 章核对数据、第 3 章起训练、第 4 章调数据管线、第 5 章排掉常见的坑半天时间就能跑出第一个能用的模型。我从一开始直接拿image_dataset_from_directory配cache和prefetch起步中间吃了几次内存和显存的亏才总结出上面这套顺序。按这个流程走希望帮到你。本文还有配套的精品资源点击获取