如何基于Keras实现从磁盘批量加载图像并完成数据增强训练？

嘿，我之前刚好遇到过和你一模一样的问题——4GB内存塞不下大训练集，还要用Keras做数据增强来平衡类别。幸运的是Keras自带了完美解决这个场景的工具，完全不用从零写自定义生成器，我一步步教你怎么弄：

这个组合既能批量从磁盘加载数据（不用一次性读入内存），又能实时对数据做增强，还能自动处理类别标签，完美匹配你的需求。

步骤1：导入必要模块 #

先把需要的库导进来：

from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense

步骤2：配置数据增强规则 #

先定义ImageDataGenerator，把你需要的增强操作加进去，同时别忘了对像素值做归一化（这是CNN训练的常规操作）：

# 训练集：开启数据增强+归一化
train_datagen = ImageDataGenerator(
    rescale=1./255,  # 将像素值缩到0-1区间
    rotation_range=40,  # 随机旋转图片（0-40度）
    width_shift_range=0.2,  # 随机水平平移图片
    height_shift_range=0.2,  # 随机垂直平移图片
    shear_range=0.2,  # 随机剪切变换
    zoom_range=0.2,  # 随机缩放图片
    horizontal_flip=True,  # 随机水平翻转图片
    fill_mode='nearest'  # 填充变换后缺失的像素
)

# 测试/验证集：只做归一化，不做增强（避免数据泄露）
test_datagen = ImageDataGenerator(rescale=1./255)

步骤3：按文件夹批量加载数据 #

首先你的数据集需要按类别分文件夹存放（这是flow_from_directory要求的标准结构）：

train/
类别A/
img1.jpg
img2.jpg
...
类别B/
img1.jpg
img2.jpg
...
validation/
类别A/
img1.jpg
...
类别B/
img1.jpg
...

然后生成数据流，每次只加载一个batch的图片，同时实时应用增强：

# 训练集数据流
train_generator = train_datagen.flow_from_directory(
    '你的训练集文件夹路径',
    target_size=(150, 150),  # 把所有图片统一resize到这个尺寸（根据你的模型调整）
    batch_size=32,  # 每次加载32张图，内存不够的话可以改成16
    class_mode='categorical'  # 多分类用categorical，二分类用binary
)

# 验证集数据流
validation_generator = test_datagen.flow_from_directory(
    '你的验证集文件夹路径',
    target_size=(150, 150),
    batch_size=32,
    class_mode='categorical'
)

步骤4：用生成器训练模型 #

在新版Keras（TensorFlow 2.x集成的版本）里，model.fit已经支持直接传入生成器了，不用再用旧版的fit_generator：

# 先定义你的CNN模型（这里举个简单的示例，你可以换成自己的模型）
model = Sequential()
model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)))
model.add(MaxPooling2D((2, 2)))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D((2, 2)))
model.add(Conv2D(128, (3, 3), activation='relu'))
model.add(MaxPooling2D((2, 2)))
model.add(Flatten())
model.add(Dense(512, activation='relu'))
model.add(Dense(2, activation='softmax'))  # 假设是2分类任务

# 编译模型
model.compile(loss='categorical_crossentropy',
              optimizer='rmsprop',
              metrics=['accuracy'])

# 开始训练，用生成器喂数据
history = model.fit(
    train_generator,
    steps_per_epoch=train_generator.samples // train_generator.batch_size,  # 每个epoch要跑多少个batch
    epochs=30,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples // validation_generator.batch_size
)

额外补充：解决类别不平衡的进阶技巧 #

除了数据增强，你还可以给样本少的类别设置更高的权重，进一步平衡训练：

# 假设类别A样本多，权重设1.0；类别B样本少，权重设5.0
class_weights = {'类别A': 1.0, '类别B': 5.0}

# 在训练时传入class_weight参数
history = model.fit(
    train_generator,
    steps_per_epoch=train_generator.samples // train_generator.batch_size,
    epochs=30,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples // validation_generator.batch_size,
    class_weight=class_weights
)

如果你的数据集不是按文件夹分类（比如标签存在CSV里），可以用flow_from_dataframe方法，用法和flow_from_directory几乎一致，只需要传入CSV路径和对应的列名即可。

这样就能同时解决内存不足和类别不平衡的问题啦！

内容的提问来源于stack exchange，提问作者Learning is a mess