使用Keras ImageDataGenerator做数据增强训练过慢，求性能优化方案

架构老林

2026-5-22

如何提升ImageDataGenerator数据增强后的训练速度？

我太懂这种感受了——本来训练10万条数据只需要3分钟，加了数据增强直接飙到3小时，简直让人崩溃。咱们先拆解下问题根源，再给你几个实用的优化方案：

首先，你的ImageDataGenerator设置里有几个拖慢速度的关键点：featurewise_center和featurewise_std_normalization会要求计算全量数据集的均值和标准差，这在10万条数据上是非常耗时的；另外实时的翻转、缩放等操作都是CPU端执行的，如果你的CPU性能一般，就会成为训练的瓶颈。

下面是具体的优化手段，按优先级排序：

1. 换掉或预计算featurewise归一化，省掉最大的耗时项

featurewise_center和featurewise_std_normalization这两个参数是最大的性能杀手——如果不提前计算，每次训练都要遍历全量数据来算均值和标准差，这直接拖慢了整个流程。

解决方法二选一：

方案A（推荐）：用简单归一化替代：把这两个参数删掉，换成rescale=1./255，直接将像素值缩放到0-1之间，效果差不多，但完全不需要计算全局统计量，速度提升非常明显。

方案B：预计算统计量：如果你一定要用featurewise归一化，先一次性计算好均值和标准差，之后训练就不用重复计算了：

# 只运行一次，计算全量训练数据的均值和标准差
datagen = ImageDataGenerator(featurewise_center=True, featurewise_std_normalization=True)
datagen.fit(train_images)  # 这里传入你的训练数据数组
# 之后训练时直接用这个datagen，它会复用预计算好的值

2. 精简数据增强操作，减少CPU计算量

不是所有增强操作都有必要，你可以评估下哪些对任务帮助大，砍掉没用的：

如果你的任务场景里垂直翻转没有意义（比如识别人脸、手写数字），直接关掉vertical_flip=True
缩小zoom_range的范围，或者直接去掉——缩放操作的计算成本不低，如果带来的精度提升有限，不如舍弃
尽量避免同时开启3种以上的像素级增强（比如翻转+缩放+旋转+平移），选2-3个最有效的就行

3. 开启多进程并行生成数据

ImageDataGenerator的flow()或flow_from_directory()支持多进程并行处理数据，把CPU的多核利用起来：

train_generator = datagen.flow(
    train_images,
    train_labels,
    batch_size=64,
    workers=4,  # 根据你的CPU核心数设置，比如8核就设4或8
    use_multiprocessing=True
)

注意：如果你的数据是存在内存里的numpy数组，多进程可能会有数据拷贝的开销，这时候用flow_from_directory直接从磁盘读取数据，配合多进程的效率会更高。

4. 离线预生成增强数据（磁盘空间足够的话）

如果你的磁盘有足够空间，可以提前把所有增强后的数据生成好，保存到磁盘上，训练时直接读取现成的文件：

用datagen.flow()生成增强数据，循环把每个batch的图片保存到指定文件夹：

import os
save_dir = "./augmented_data"
os.makedirs(save_dir, exist_ok=True)
generator = datagen.flow(train_images, train_labels, batch_size=32, save_to_dir=save_dir, save_prefix="aug", save_format="jpg")
# 生成足够的增强数据，比如生成和原数据量相同的增强数据
for i in range(len(train_images)//32):
    generator.next()

训练时用flow_from_directory直接读取这些预生成的文件，速度会和你没开增强时差不多——相当于用磁盘空间换时间。

5. 切换到tf.data.Dataset（进阶优化）

ImageDataGenerator其实是比较老的API了，TensorFlow的tf.data.Dataset配合tf.image里的增强操作，能把部分计算转移到GPU上，效率更高：

import tensorflow as tf

def augment_image(image, label):
    image = tf.image.random_flip_left_right(image)
    # 按需添加其他增强操作，比如随机缩放
    # image = tf.image.random_zoom(image, [0.8, 1.2])
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

# 构建数据集
train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_dataset = train_dataset.shuffle(len(train_images)).batch(64).map(augment_image, num_parallel_calls=tf.data.AUTOTUNE)

# 训练模型
model.fit(train_dataset, epochs=10)

这个方法的优势是增强操作可以和模型训练并行，甚至部分操作在GPU上执行，速度会比ImageDataGenerator快很多。

6. 硬件和训练参数微调

确保你的磁盘是SSD：机械硬盘读取10万条数据的速度本身就慢，再加上实时增强，会雪上加霜；SSD能大幅提升数据读取速度。
适当增大batch size：更大的batch size能减少迭代次数，同时让数据生成的效率更高（进程切换的开销更小）。

最后给你一个优化后的完整代码示例参考：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 优化后的ImageDataGenerator设置
datagen = ImageDataGenerator(
    horizontal_flip=True,
    # vertical_flip=True,  # 按需保留
    rescale=1./255  # 用简单归一化替代featurewise操作
)

# 开启多进程生成数据
train_generator = datagen.flow(
    train_images,
    train_labels,
    batch_size=64,
    workers=4,
    use_multiprocessing=True
)

# 训练模型
model.fit(
    train_generator,
    steps_per_epoch=len(train_images) // 64,
    epochs=10
)

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

使用Keras ImageDataGenerator做数据增强训练过慢，求性能优化方案

1. 换掉或预计算featurewise归一化，省掉最大的耗时项 #

2. 精简数据增强操作，减少CPU计算量 #

3. 开启多进程并行生成数据 #

4. 离线预生成增强数据（磁盘空间足够的话） #

5. 切换到tf.data.Dataset（进阶优化） #