TensorFlow2.X小数据集图像分类：MobileNetV2迁移学习实战

简介这份资源面向希望上手深度学习图像分类的开发者与学习者聚焦TensorFlow 2.X环境下MobileNetV2模型的实战应用。内容基于植物幼苗数据集中的部分样本覆盖12个类别适合作为小数据集迁移学习的练手项目。压缩包共约2000个文件以png图片数据为主另含4个Python脚本、1份PDF说明文档和1个h5模型文件整体约961.42MB图片用于训练与验证脚本负责数据加载、标签onehot编码、数据增强、mixup、数据集切分及预训练模型加载等环节。目前已有753人学习下载。通过这份资源读者可以完整走通从数据准备到模型加载的流程理解线性瓶颈与倒残差结构在轻量级网络中的作用并借助现成脚本与模型文件快速复现实验、对照排查问题为移动端图像分类任务打下实践基础。1. 小数据集也能跑 MobileNetV2这份 TensorFlow2.X 图像分类包到底给了什么手里只有几百张图却要做一个 12 类别的图像分类任务这种场景在工业质检、农业识别、医疗辅诊里太常见了。很多人第一反应是上 ResNet50 或者 EfficientNet结果训练集准确率冲到 99%验证集死活上不去典型的过拟合翻车现场。这份资源给了一条更务实的路线用 TensorFlow2.X 加载 MobileNetV2 预训练权重配合数据增强和 mixup在植物幼苗小数据集上做 12 分类。MobileNetV2 的核心是倒残差结构和线性瓶颈参数量只有 3.4M 左右推理速度快适合往移动端或者边缘设备上落。资源包里包含my_model.h5训练好的权重文件、一份 PDF 说明文档以及若干张训练过程截图。如果你手头的数据集规模不大又想快速验证一个图像分类算法能不能跑通这份东西能帮你省掉搭骨架的时间。下面我从数据加载一路讲到模型保存把每个环节的参数和坑都拆开说。2. 数据管道与标签编码从文件夹到 tf.data.Dataset 的完整链路2.1 为什么小数据集必须走 tf.data 而不是 ImageDataGeneratorTensorFlow2.X 里做图像分类常见做法有两种ImageDataGenerator的flow_from_directory和tf.data.Dataset的image_dataset_from_directory。前者是 Keras 老牌接口后者是 TF2 原生推荐。小数据集上两者都能跑但tf.data的优势在于管道可以预取、缓存、并行化而且和tf.keras的fit配合更顺。我一般会优先用image_dataset_from_directory因为它直接吃文件夹结构标签自动按子目录名生成省掉手写映射的麻烦。资源里的植物幼苗数据集按类别分文件夹存放每个子文件夹是一种幼苗。加载时指定image_size(224, 224)、batch_size32、label_modecategorical这样标签直接就是 onehot 编码不用再手动转。注意label_mode有三个可选值int返回整数标签categorical返回 onehotbinary用于二分类。12 分类任务必须用categorical否则后面算损失函数时维度对不上。import tensorflow as tf # 训练集路径和验证集路径按实际目录调整 train_dir data/train val_dir data/val train_ds tf.keras.preprocessing.image_dataset_from_directory( train_dir, image_size(224, 224), # MobileNetV2 默认输入尺寸 batch_size32, label_modecategorical, # 12 分类输出 onehot shuffleTrue, seed42 ) val_ds tf.keras.preprocessing.image_dataset_from_directory( val_dir, image_size(224, 224), batch_size32, label_modecategorical, shuffleFalse # 验证集不打乱方便对齐标签 )逻辑说明image_dataset_from_directory会扫描目录下所有子文件夹按文件夹名排序后生成类别索引。shuffleTrue只在训练集开验证集关掉否则评估时标签和预测对不上。seed固定后每次运行划分一致方便复现。参数上batch_size根据显存调8G 显存跑 224×224 的 MobileNetV232 基本安全12G 以上可以上 64。2.2 标签 onehot 与类别数校验虽然label_modecategorical已经自动转了 onehot但有一件事必须做确认类别数。资源里是 12 类但如果你换了自己的数据集类别数变了模型最后一层的Dense单元数必须跟着改。我习惯在加载后立刻打印class_names然后手动核对。class_names train_ds.class_names num_classes len(class_names) print(f类别数: {num_classes}, 类别名: {class_names}) # 检查一个 batch 的标签形状 for images, labels in train_ds.take(1): print(f图像 batch 形状: {images.shape}) # (32, 224, 224, 3) print(f标签 batch 形状: {labels.shape}) # (32, 12)如果标签形状第二维不是 12说明label_mode设错了或者子文件夹里混了非类别目录。常见坑是数据集里有个.ipynb_checkpoints或者__MACOSX文件夹被当成一个类别导致类别数变成 13。解决办法是在加载前用脚本清理非图像目录或者手动指定class_names参数。2.3 数据增强与 mixup 的接入位置小数据集上数据增强是刚需。资源里用了随机翻转、旋转、缩放、对比度调整这一套。TF2 里可以用tf.keras.layers.RandomFlip、RandomRotation、RandomZoom这些预处理层直接塞进模型前面或者放在tf.data管道里用.map()做。我一般放在模型里因为这样保存h5时增强逻辑一起带走推理时自动关闭。mixup 稍微特殊一点它不是单张图变换而是把两张图按比例混合标签也按同样比例混合。实现上要在 batch 级别操作所以得用tf.data的.map()在 batch 之后做。import tensorflow as tf def mixup(images, labels, alpha0.2): batch_size tf.shape(images)[0] # 从 Beta 分布采样混合系数 lam tf.compat.v1.distributions.Beta(alpha, alpha).sample() # 打乱索引 indices tf.random.shuffle(tf.range(batch_size)) shuffled_images tf.gather(images, indices) shuffled_labels tf.gather(labels, indices) # 混合 mixed_images lam * images (1 - lam) * shuffled_images mixed_labels lam * labels (1 - lam) * shuffled_labels return mixed_images, mixed_labels train_ds train_ds.map(mixup, num_parallel_callstf.data.AUTOTUNE) train_ds train_ds.prefetch(tf.data.AUTOTUNE)参数说明alpha0.2是 mixup 的强度越小混合越接近原图越大越模糊。小数据集上我一般用 0.1 到 0.3太大反而欠拟合。num_parallel_calls设AUTOTUNE让 TF 自己决定并行数prefetch提前取下一批数据减少 GPU 等待。注意 mixup 之后标签不再是严格的 onehot而是浮点混合值所以损失函数要用CategoricalCrossentropy它支持软标签。3. MobileNetV2 迁移学习冻结策略、学习率与模型保存3.1 加载预训练权重与冻结层数选择MobileNetV2 在tf.keras.applications里直接可用weightsimagenet加载预训练权重include_topFalse去掉原来的 1000 类分类头。小数据集上迁移学习的标准做法是先冻结骨干网络只训练新加的分类头训几个 epoch 后再解冻一部分高层做微调。from tensorflow.keras import layers, models from tensorflow.keras.applications import MobileNetV2 base_model MobileNetV2( input_shape(224, 224, 3), include_topFalse, weightsimagenet ) # 先全部冻结 base_model.trainable False # 构建分类头 model models.Sequential([ base_model, layers.GlobalAveragePooling2D(), layers.Dropout(0.3), layers.Dense(num_classes, activationsoftmax) ]) model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-3), losscategorical_crossentropy, metrics[accuracy] ) model.summary()逻辑说明GlobalAveragePooling2D把特征图从(7, 7, 1280)压成(1280,)比Flatten参数少得多不容易过拟合。Dropout(0.3)在小数据集上是保险丝比例再高可能欠拟合。学习率1e-3是 Adam 的常用起点冻结阶段可以稍大微调阶段必须降到1e-5量级否则预训练权重会被冲垮。3.2 微调阶段的解冻与学习率重设冻结训练 10 到 15 个 epoch 后验证集准确率一般能到 80% 以上。这时候解冻 MobileNetV2 最后 30 到 50 层做微调学习率降到1e-5。解冻层数不是越多越好小数据集上解冻太多照样过拟合。# 解冻最后 40 层 base_model.trainable True for layer in base_model.layers[:-40]: layer.trainable False # 重新编译学习率调低 model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-5), losscategorical_crossentropy, metrics[accuracy] ) # 继续训练 history_fine model.fit( train_ds, validation_dataval_ds, epochs20, callbacks[ tf.keras.callbacks.EarlyStopping(patience5, restore_best_weightsTrue), tf.keras.callbacks.ModelCheckpoint(best_model.h5, save_best_onlyTrue) ] )参数说明EarlyStopping的patience5表示验证集损失连续 5 个 epoch 不降就停restore_best_weightsTrue回滚到最优权重这是后悔药级别的配置。ModelCheckpoint只存最优模型避免最后 epoch 过拟合的权重被保存。注意重新编译是必须的否则学习率改动不生效。3.3 模型保存与 h5 格式的注意事项资源里给的my_model.h5就是这种保存方式的产物。h5格式在 TF2.X 里仍然支持但官方更推荐SavedModel或者.keras格式。h5的问题是自定义层和 mixup 这类非标准操作可能存不全。如果模型里用了tf.keras.layers.RandomFlip这些预处理层h5能存但加载时需要确保 TF 版本一致。# 保存完整模型 model.save(my_model.h5) # 加载时 loaded_model tf.keras.models.load_model(my_model.h5) loaded_model.evaluate(val_ds)如果加载时报Unknown layer或者Lambda相关错误说明模型里有自定义函数。解决办法是用custom_objects参数传入或者改用SavedModel格式。我一般会在保存前跑一遍model.evaluate确认推理正常再存避免存了个坏模型还不知道。4. 避坑与排查小数据集训练 MobileNetV2 的五个血泪经验4.1 验证集准确率震荡大时高时低现象每个 epoch 验证集准确率跳动超过 10%损失曲线锯齿状。原因通常有两个一是验证集太小二三十张图一个 batch 的波动就能让指标大幅摆动二是shuffleFalse没设验证集顺序固定但模型预测不稳定。解决验证集至少每类 10 到 15 张总量不低于 100 张验证集shuffle关掉但评估时用model.evaluate而不是手动循环加BatchNormalization的模型在推理时要确保trainingFalseevaluate会自动处理。4.2 mixup 之后损失变成 NaN现象训练几个 step 后 loss 直接 NaN。原因mixup 的lam采样用了tf.compat.v1.distributions.Beta在某些 TF 版本里返回的是标量张量和 batch 维度广播时出错或者alpha设得太大混合后标签值过小CategoricalCrossentropy的from_logits参数没设对。解决确认from_logitsFalse因为模型最后一层是softmaxalpha控制在 0.2 以内用tf.clip_by_value把混合后的标签裁剪到[1e-7, 1-1e-7]。4.3 冻结训练时准确率不涨现象前 10 个 epoch 准确率卡在 10% 左右跟随机猜差不多。原因分类头初始化用了默认的glorot_uniform但Dense层前面接了GlobalAveragePooling2D特征值范围偏小梯度传不回去。解决把分类头的Dense初始化改成he_normal或者在GlobalAveragePooling2D后面加一个BatchNormalization。另一个常见原因是学习率太低冻结阶段用1e-3而不是1e-4。4.4 解冻微调后验证集准确率反而下降现象冻结阶段验证集到 85%解冻后掉到 70%。原因解冻层数太多或者学习率没降。MobileNetV2 的浅层学的是通用边缘纹理小数据集上微调这些层等于破坏预训练特征。解决只解冻最后 20 到 30 层学习率降到1e-5甚至1e-6并且解冻后前两个 epoch 用 warmup 慢慢升学习率。4.5 h5 模型加载后预测结果全为同一类现象训练时验证集正常保存后重新加载预测所有图都输出同一个类别。原因h5保存时没存优化器状态但这不影响推理真正的问题通常是加载时compileFalse导致某些自定义层没初始化或者输入图像的预处理和训练时不一致。MobileNetV2 的preprocess_input在训练时如果用了推理时也必须用否则输入分布偏移。解决把预处理逻辑写进模型里用tf.keras.layers.Rescaling或者Lambda层包住保存时一起带走。5. 从 h5 到实际推理单张图预测与批量评估的落地技巧训练完拿到my_model.h5只是第一步真正要用起来得会做单张推理和批量评估。单张图预测的坑在于输入维度model.predict接受的是(batch, 224, 224, 3)直接传一张(224, 224, 3)的图会报维度错误。我一般用tf.expand_dims加一个 batch 维。import numpy as np from tensorflow.keras.preprocessing import image def predict_single(img_path, model, class_names): img image.load_img(img_path, target_size(224, 224)) x image.img_to_array(img) / 255.0 # 归一化到 [0,1] x np.expand_dims(x, axis0) # (1, 224, 224, 3) preds model.predict(x, verbose0) idx np.argmax(preds[0]) return class_names[idx], preds[0][idx] label, score predict_single(test.jpg, loaded_model, class_names) print(f预测: {label}, 置信度: {score:.4f})注意归一化方式必须和训练时一致。如果训练用了preprocess_input推理也得用不能简单除以 255。批量评估更简单直接model.evaluate(val_ds)拿准确率但要看混淆矩阵的话得手动跑预测。from sklearn.metrics import confusion_matrix import seaborn as sns import matplotlib.pyplot as plt y_true [] y_pred [] for images, labels in val_ds: preds loaded_model.predict(images, verbose0) y_true.extend(np.argmax(labels.numpy(), axis1)) y_pred.extend(np.argmax(preds, axis1)) cm confusion_matrix(y_true, y_pred) sns.heatmap(cm, annotTrue, fmtd, xticklabelsclass_names, yticklabelsclass_names) plt.show()混淆矩阵能看出哪些类别容易混。植物幼苗数据集里不同种类的幼苗在早期形态上非常接近混淆矩阵上出现 20% 以上的误判很正常。这时候可以考虑加类别权重或者对易混类别做针对性增强。我自己的习惯是每次训完模型先跑一遍混淆矩阵再决定要不要调参。从那以后我每次保存 h5 之前都强制走一遍单张推理和混淆矩阵确认模型不是个黑匣子才敢往生产环境放。希望帮到你。本文还有配套的精品资源点击获取