如何在TensorFlow中基于大量图片创建适配MobileNetV2的数据集

第一步：修正现有配置错误 #

• 测试集路径配置错误：你当前的test_path指向了训练集目录，需要调整为测试集对应路径
• 类别名大小写不匹配：你定义的categories是小写开头，实际目录名是首字母大写的Freshapple等，需保持二者一致，修正后配置如下：

train_path = ".../Desktop/Data/fruit-dataset/train"
test_path = ".../Desktop/Data/fruit-dataset/test"
categories = ["Freshapple", "Freshorange", "Freshbanana",
              "Rottenapple", "Rottenorange", "Rottenbanana"]

第二步：配置数据生成器 #

适配MobilenetV2的预处理要求，同时给训练集增加数据增强提升泛化能力，测试集仅做标准化处理：

import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 训练集生成器：带数据增强
train_datagen = ImageDataGenerator(
    preprocessing_function=tf.keras.applications.mobilenet_v2.preprocess_input,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True
)

# 测试集生成器：仅做预处理，不做增强
test_datagen = ImageDataGenerator(
    preprocessing_function=tf.keras.applications.mobilenet_v2.preprocess_input
)

第三步：调用flow_from_directory流式加载数据集 #

采用批量流式读取逻辑，不会全量加载图片到内存，彻底解决13000张图内存崩溃问题：

# 加载训练集
train_generator = train_datagen.flow_from_directory(
    directory=train_path,
    target_size=(224, 224), # 匹配MobilenetV2默认输入尺寸
    batch_size=32, # 可根据显存大小调整，显存小可以设为16或者8
    class_mode="categorical", # 6分类多分类任务用categorical
    classes=categories,
    shuffle=True
)

# 加载测试集
test_generator = test_datagen.flow_from_directory(
    directory=test_path,
    target_size=(224, 224),
    batch_size=32,
    class_mode="categorical",
    classes=categories,
    shuffle=False # 测试集不需要打乱，方便后续评估结果对齐
)

后续使用说明 #

生成的train_generator和test_generator可以直接传入model.fit()用于训练，层冻结、微调等实验不需要修改数据加载逻辑。如果需要更高的IO训练效率，可以替换为tf.keras.utils.image_dataset_from_directory()接口，用法逻辑基本一致，底层采用tf.data流水线优化，训练速度更快。

内容的提问来源于stack exchange，提问作者LeGOATJames23