TensorFlow实战：8万张图245类垃圾分类模型训练与调优

简介这份资源面向深度学习入门者与计算机视觉实践者提供一套完整的垃圾分类图像分类方案包含约8万张图片、覆盖245个类别的数据集以及基于TensorFlow实现的训练与推理代码。压缩包共73个文件约101.99MB其中jpg、png、jpeg等图片文件构成数据集主体py脚本负责模型训练与测试h5文件为已训练好的模型权重另有txt、xml、md等配置与说明文档整体结构清晰便于直接上手复现。资源在models目录下附带两组训练完成的模型可省去从零训练的时间成本适合用于课程设计、毕业项目或算法对比实验。代码部分涵盖CNN与MobileNet等不同网络结构的训练脚本并配有工具函数与测试入口方便读者理解数据加载、模型构建、训练调参到评估的完整流程。目前已有29958人学习下载可作为垃圾分类任务中数据准备、模型选型与工程落地的参考范例。1. 8 万张图、245 个类垃圾分类数据集到底能拿来做什么手里有一份 8 万张图片、245 个类别的垃圾分类数据集还配了 TensorFlow 代码这件事本身就值得认真对待。多数公开垃圾分类数据集要么只有 6 个大类可回收、厨余、有害、其他这种粗分要么样本量撑死一两万张训出来的模型在真实场景里一遇到「奶茶杯算哪类」就翻车。245 个类意味着它已经细到具体物品级别——纸杯、塑料瓶、玻璃瓶、电池、灯管、快递纸箱这些各自成类8 万张的体量也够撑起一个从零训练的卷积网络或者微调一个预训练骨干。这份资源适合三类人想做一个能真正落地的垃圾分类识别 demo 的开发者、需要多类别细粒度分类数据集做课程设计或论文实验的学生、以及想拿它当 backbone 验证自己训练 pipeline 的算法工程师。下面我按「数据长什么样 → 怎么读进来 → 怎么训 → 坑在哪 → 怎么调优」的顺序把这条链路走一遍。2. 先搞清楚数据长什么样目录结构、类别分布与 TFRecord 取舍拿到一个 zip 包第一件事不是急着写模型而是把数据摸清楚。245 个类、8 万张图平均每类 300 多张但真实分布几乎不可能均匀——有些类可能上千张有些冷门类可能只有几十张。这个不均衡直接决定你后面要不要做重采样、要不要用 class weight、评估时看 accuracy 还是看 macro-F1。2.1 目录结构与类别映射的读取方式常见的组织方式是每个类别一个子目录目录名就是类别名图片放在里面。先用一段脚本把结构、类别数、每类数量统计出来别凭感觉。import os import json from collections import Counter DATA_ROOT ./garbage_dataset # 解压后的根目录按实际改 IMG_EXT (.jpg, .jpeg, .png, .bmp) # 遍历一级子目录把目录名当类别名 class_names sorted([ d for d in os.listdir(DATA_ROOT) if os.path.isdir(os.path.join(DATA_ROOT, d)) ]) print(类别总数:, len(class_names)) counter Counter() bad_files [] for cls in class_names: cls_dir os.path.join(DATA_ROOT, cls) for f in os.listdir(cls_dir): if f.lower().endswith(IMG_EXT): counter[cls] 1 else: bad_files.append(os.path.join(cls_dir, f)) print(图片总数:, sum(counter.values())) print(最多类:, counter.most_common(3)) print(最少类:, counter.most_common()[-3:]) print(非图片文件数:, len(bad_files)) # 保存类别到索引的映射训练和推理必须用同一份 with open(class_index.json, w, encodingutf-8) as fp: json.dump({name: i for i, name in enumerate(class_names)}, fp, ensure_asciiFalse, indent2)这段脚本做了三件事确定类别顺序用sorted保证可复现别用os.listdir的原始顺序那个顺序在不同机器上可能不一样、统计每类样本量、把非图片文件挑出来。class_index.json是关键训练时类别索引和推理时类别索引必须来自同一份映射否则模型输出第 37 类你查表查成另一个东西这种 bug 排查起来非常费劲。参数上IMG_EXT按你实际数据补有些数据集混了.webp或.tif漏掉就会少统计。2.2 类别不均衡先量化再决定要不要处理统计完你会看到分布。如果最多类和最少类差 10 倍以上直接上tf.data的sample_from_datasets或者给class_weight赋值。差 3 倍以内通常不用特殊处理靠数据增强和足够的 epoch 就能压过去。这里给一个按类别加权的思路权重取样本数的倒数再归一化import numpy as np counts np.array([counter[c] for c in class_names], dtypenp.float32) class_weight counts.sum() / (len(counts) * counts) # 反比权重 class_weight class_weight / class_weight.mean() # 归一化均值拉到 1 for c, w in zip(class_names[:5], class_weight[:5]): print(c, round(float(w), 3))class_weight在model.fit里以字典形式传入键是类别索引。注意权重别拉太猛反比权重在极端不均衡时会让少数类权重爆炸训练震荡我一般会开方缓一下class_weight np.sqrt(class_weight)。2.3 用 tf.data 搭输入管道别用 ImageDataGeneratorTensorFlow 现在做图像输入tf.keras.utils.image_dataset_from_directory或者自己写tf.data管道都比老的ImageDataGenerator稳。8 万张图如果全读进内存按 224×224×3 的 uint8 算大概 12GB普通机器扛不住所以必须走流式加载。import tensorflow as tf IMG_SIZE (224, 224) BATCH 64 AUTOTUNE tf.data.AUTOTUNE train_ds tf.keras.utils.image_dataset_from_directory( DATA_ROOT, validation_split0.2, subsettraining, seed42, image_sizeIMG_SIZE, batch_sizeBATCH, label_modeint, ) val_ds tf.keras.utils.image_dataset_from_directory( DATA_ROOT, validation_split0.2, subsetvalidation, seed42, image_sizeIMG_SIZE, batch_sizeBATCH, label_modeint, ) # 缓存 预取避免 GPU 等数据 train_ds train_ds.cache().shuffle(2000).prefetch(AUTOTUNE) val_ds val_ds.cache().prefetch(AUTOTUNE)seed42保证训练集和验证集划分可复现label_modeint配合 sparse 交叉熵损失。cache()在数据能放进内存时加速明显放不下就换成cache(./cache_train)落盘或者干脆去掉。shuffle(2000)的缓冲区别设太小245 个类如果缓冲区只有几百一个 batch 里可能全是同一类BatchNorm 的统计量会偏。提示image_dataset_from_directory默认按目录名字母序生成类别索引和你自己sorted出来的顺序一致但如果你手动改过目录名或用了自定义映射务必核对train_ds.class_names和class_index.json是否对得上。3. 用 TensorFlow 把 245 类分类模型训起来骨干选择与两阶段策略245 个类不算少从零训一个 ResNet 也能收敛但更省算力、更容易出效果的做法是拿 ImageNet 预训练骨干做微调。TensorFlow 里tf.keras.applications提供了现成的骨干选哪个、怎么接分类头、冻结多少层直接决定你第一版模型能不能用。3.1 骨干选型EfficientNetV2 还是 ResNet50我的默认选择是EfficientNetV2B0参数量小、在 224 分辨率下精度和速度平衡好适合 8 万张这个量级。如果你机器显存充裕、追求更高上限ResNet50或EfficientNetV2B3也行但训练时间会明显拉长。分类头用GlobalAveragePooling2D加一层 Dropout 再接到 245 维输出别用 Flatten参数量会炸。from tensorflow.keras import layers, models NUM_CLASSES len(class_names) base tf.keras.applications.EfficientNetV2B0( include_topFalse, weightsimagenet, input_shape(224, 224, 3) ) base.trainable False # 第一阶段冻结骨干 inputs layers.Input(shape(224, 224, 3)) # EfficientNetV2 自带预处理这里不要再手动 rescale x base(inputs, trainingFalse) x layers.GlobalAveragePooling2D()(x) x layers.Dropout(0.3)(x) outputs layers.Dense(NUM_CLASSES, activationsoftmax)(x) model models.Model(inputs, outputs) model.compile( optimizertf.keras.optimizers.Adam(1e-3), losstf.keras.losses.SparseCategoricalCrossentropy(), metrics[accuracy, tf.keras.metrics.SparseTopKCategoricalAccuracy(k5)], ) model.summary()关键点EfficientNetV2系列在 Keras 里已经内置了归一化你不要再套一层Rescaling(1./255)否则输入分布错了精度会莫名其妙掉一截这是很常见的翻车点。SparseTopKCategoricalAccuracy(k5)对 245 类特别有用top-1 可能只有六七成但 top-5 往往能到九成以上能帮你看清模型到底学没学到东西。3.2 两阶段训练先冻结后解冻学习率要分开设第一阶段冻结骨干只训分类头通常 5 到 10 个 epoch 就能让 loss 明显下降。第二阶段解冻骨干的后半部分做微调学习率必须调小否则预训练权重会被冲掉。callbacks [ tf.keras.callbacks.ModelCheckpoint( best_head.keras, monitorval_accuracy, save_best_onlyTrue), tf.keras.callbacks.EarlyStopping( monitorval_accuracy, patience4, restore_best_weightsTrue), ] # 阶段一只训分类头 history1 model.fit( train_ds, validation_dataval_ds, epochs8, callbackscallbacks, class_weightdict(enumerate(class_weight)), ) # 阶段二解冻骨干后 40% 的层做微调 base.trainable True fine_tune_at int(len(base.layers) * 0.6) for layer in base.layers[:fine_tune_at]: layer.trainable False model.compile( optimizertf.keras.optimizers.Adam(1e-5), # 微调学习率降两个数量级 losstf.keras.losses.SparseCategoricalCrossentropy(), metrics[accuracy, tf.keras.metrics.SparseTopKCategoricalAccuracy(k5)], ) history2 model.fit( train_ds, validation_dataval_ds, epochs20, callbackscallbacks, class_weightdict(enumerate(class_weight)), )fine_tune_at取 0.6 是个经验值意思是只动骨干靠后的高层特征前面的底层边缘、纹理特征保持不动。如果你的数据域和 ImageNet 差得远比如全是俯拍的垃圾桶内部可以把这个比例降到 0.4多解冻一些。class_weight用dict(enumerate(...))转成索引到权重的映射和label_modeint对应。3.3 数据增强放在模型里还是管道里TensorFlow 现在推荐把增强层直接写进模型这样导出 SavedModel 时增强逻辑跟着走推理时自动关闭。用RandomFlip、RandomRotation、RandomZoom这几层就够别上太猛的增强垃圾分类的物体形状和颜色是判别关键过度裁剪或颜色抖动反而伤精度。data_aug tf.keras.Sequential([ layers.RandomFlip(horizontal), layers.RandomRotation(0.08), layers.RandomZoom(0.1), layers.RandomContrast(0.1), ], nameaug) # 在模型里base 之前插入 x data_aug(inputs) x base(x, trainingFalse)RandomRotation(0.08)是旋转 ±8% 的 2π约 ±29 度对垃圾图片够用了。RandomContrast(0.1)模拟不同光照真实场景里垃圾桶光线差异很大这一层能提升泛化。注意增强层只在trainingTrue时生效验证和推理自动跳过不用你手动切。4. 训练过程中的避坑清单从 OOM 到精度不涨的排查这一节是我踩过的坑里最费时间的几个按「现象 → 原因 → 解决」写你遇到类似症状可以直接对号入座。4.1 显存爆了OOM不一定是 batch 太大现象训练几个 step 后报ResourceExhaustedError或者进程被系统 kill。原因常见的是cache()把整个数据集缓存进内存后又复制到显存或者shuffle缓冲区设得过大。另一个隐蔽原因是验证集也开了cache()两份缓存叠加。解决先去掉cache()或改成落盘缓存shuffle缓冲区降到 1000 以内把 batch 从 64 降到 32 并配合梯度累积。用tf.config.experimental.set_memory_growth打开显存按需增长避免 TF 一上来就占满。4.2 训练精度一直卡在 1/245 附近现象loss 不降accuracy 约等于随机猜。原因多半是标签和输入没对齐或者预处理重复做了。比如骨干自带归一化你又套了一层Rescaling输入全被压到接近 0网络学不到东西。解决打印一个 batch 的图片像素范围确认在合理区间EfficientNetV2 期望 0-255ResNet 期望经过preprocess_input。再核对class_index.json和train_ds.class_names是否一致。这两个查完基本能定位。4.3 验证集精度远高于训练集现象val_accuracy 比 train_accuracy 高十几个点。原因Dropout 和增强只在训练时开验证时关了如果训练集本身难、验证集恰好简单就会出现这种倒挂。另一个原因是验证集太小波动大。解决先确认验证集划分是否随机、量是否够至少每类几张。如果确实是增强太猛导致训练集「看起来难」把RandomRotation、RandomZoom的幅度调小。别急着高兴这种倒挂往往意味着模型没学到判别性特征。4.4 微调后精度反而掉了现象阶段一 val_accuracy 到 0.7阶段二解冻后掉到 0.5。原因微调学习率太大把预训练权重冲毁了或者解冻层数太多底层特征被破坏。解决学习率降到 1e-5 甚至 1e-6fine_tune_at往上调少解冻一些并且微调阶段前两个 epoch 用 warmup 让学习率从 0 线性升上去。我一般会在微调前先跑一个 epoch 只观察 loss确认没炸再继续。4.5 保存的模型推理时类别对不上现象训练时 accuracy 正常加载模型预测输出的类别名全是错的。原因推理脚本里的类别列表和训练时的顺序不一致或者用了os.listdir的原始顺序。解决推理时强制从class_index.json读映射并且把这份 json 和模型一起保存。导出 SavedModel 时可以把类别名写进模型的signatures但最稳的还是外部维护一份映射文件加载模型后立刻校验类别数是否等于 245。5. 把模型用起来导出、量化与一个提升 top-1 的实用技巧训完模型只是半程真正落地还要考虑导出格式、推理速度和实际场景的鲁棒性。这一节讲怎么把 245 类模型变成一个能塞进服务或边缘设备的产物以及一个我常用的、能实打实提点的技巧。5.1 导出 SavedModel 与 TFLite 量化服务端部署用 SavedModel 最省事边缘设备或移动端用 TFLite。245 类的全连接层输出不大量化收益主要在骨干部分。# 导出 SavedModel model.save(garbage_245_savedmodel) # 转 TFLite 并做动态范围量化 converter tf.lite.TFLiteConverter.from_saved_model(garbage_245_savedmodel) converter.optimizations [tf.lite.Optimize.DEFAULT] tflite_model converter.convert() with open(garbage_245.tflite, wb) as f: f.write(tflite_model) print(TFLite 大小(MB):, round(len(tflite_model) / 1024 / 1024, 2))动态范围量化只量化权重激活值保持浮点精度损失通常 1 个点以内模型体积能压到原来的四分之一左右。如果你要上 INT8 全整型量化需要准备代表性数据集做校准精度可能再掉 1 到 2 个点245 类这种细粒度任务要谨慎建议先跑一版对比 top-5 再决定。5.2 用测试时增强TTA把 top-1 再抬一抬细粒度分类里单张图预测容易受拍摄角度影响。测试时增强的做法是对同一张图做几种变换把预测概率平均。代价是推理时间翻几倍但对精度敏感的场景很值。import numpy as np def predict_with_tta(model, img_batch): # img_batch: (N, 224, 224, 3)已归一化到模型期望范围 probs model.predict(img_batch, verbose0) probs model.predict(img_batch[:, :, ::-1, :], verbose0) # 水平翻转 probs model.predict(np.rot90(img_batch, 1, axes(1, 2)), verbose0) # 旋转90度 return probs / 3.0 # 用法取一个 batch 验证 for imgs, labels in val_ds.take(1): avg_probs predict_with_tta(model, imgs.numpy()) preds np.argmax(avg_probs, axis1) acc np.mean(preds labels.numpy()) print(TTA top-1:, round(float(acc), 4))三个视角平均是最小可用的 TTA翻转和旋转各一次。注意np.rot90之后图片尺寸还是 224×224不用 resize。如果你的模型对旋转敏感比如训练时没加旋转增强TTA 提升会很明显如果训练时已经加了强旋转增强TTA 收益可能只有零点几个点那就没必要上省推理时间。5.3 一个我常犯的错忽略类别名的可读性最后说个不是技术但很要命的点。245 个类的目录名如果是拼音缩写或者编号比如class_037你训完模型自己都不知道第 37 类是什么。我现在的习惯是拿到数据集第一件事就是把类别名整理成中文可读的映射表和class_index.json一起维护。推理输出时直接显示「塑料瓶」而不是「37」这个习惯帮我省了无数次回头翻目录的时间。模型精度可以慢慢调但类别映射错了整个系统就是黑匣子连后悔药都没得吃。希望这些步骤和踩坑记录能帮你把这份 8 万张、245 类的垃圾分类数据集真正跑起来少走点我当年走过的弯路。本文还有配套的精品资源点击获取