多次初始化Keras ImageDataGenerator出现明显内存泄漏问题

你的内存持续增长问题核心在于：每次调用populate_set都会创建新的ImageDataGenerator、大尺寸numpy数组和生成器对象，而旧对象因生成器的无限迭代特性（while True的yield）或TensorFlow内部引用，无法被Python垃圾回收机制自动清理，最终导致内存耗尽。

以下是三种针对性的解决方案：

方案1：复用生成器，仅更新数据 #

避免重复创建生成器和数据增强器，初始化一次后仅替换数据，从根源减少对象创建和内存占用。

import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator
seed=100

def create_generator(elements=4000):
    # 初始化数据增强配置（仅执行一次）
    data_gen_args = dict(
        featurewise_center=False,
        samplewise_center=False,
        featurewise_std_normalization=False,
        samplewise_std_normalization=False,
        rotation_range=7.0,
        width_shift_range=0.04,
        height_shift_range=0.04,
        shear_range=0.025,
        zoom_range=0.03,
        fill_mode='constant',
        cval=0,
        horizontal_flip=False,
        vertical_flip=False
    )
    image_datagen = ImageDataGenerator(**data_gen_args)
    mask_datagen = ImageDataGenerator(**data_gen_args)
    
    # 初始化空数据占位符
    images = np.zeros((elements, 224, 224, 6), dtype=np.uint8)
    labels = np.zeros((elements, 224, 224, 1), dtype=np.uint8)
    
    # 创建生成器
    image_generator = image_datagen.flow(images, batch_size=16, seed=seed)
    mask_generator = mask_datagen.flow(labels, batch_size=16, seed=seed)
    
    # 组合生成器
    def combine_generator():
        while True:
            yield (next(image_generator), next(mask_generator))
    
    # 定义数据更新函数
    def update_data(new_images, new_labels):
        image_generator.x = new_images
        mask_generator.x = new_labels
        # 重置生成器状态，避免残留迭代位置
        image_generator.reset()
        mask_generator.reset()
    
    return combine_generator(), update_data

# 仅初始化一次生成器和更新函数
max_elem = 4000
max_valid = 2000
train_generator, update_train_data = create_generator(elements=max_elem)
test_generator, _ = create_generator(elements=max_valid)

for _ in range(0,1000):
    # 替换为你实际加载新训练数据的逻辑
    new_train_imgs = np.zeros((max_elem, 224, 224, 6), dtype=np.uint8)
    new_train_labels = np.zeros((max_elem, 224, 224, 1), dtype=np.uint8)
    
    # 更新生成器数据
    update_train_data(new_train_imgs, new_train_labels)
    
    # 执行训练
    # hist = model.fit(train_generator, validation_data=test_generator,
    #                  validation_steps=steps_per_epoch_val,
    #                  steps_per_epoch=steps_per_epoch, epochs=epochs_together)

方案2：手动触发垃圾回收 #

在替换数据前显式删除旧对象引用，并强制触发垃圾回收，适合快速调试现有代码。

import numpy as np
import gc
from tensorflow.keras.preprocessing.image import ImageDataGenerator
seed=100

# 保留你原有的combine_generator和populate_set函数不变
def combine_generator(gen0, gen1):
    while True:
        yield (next(gen0), next(gen1))

def populate_set(elements=1e6):
    data_gen_args = dict(featurewise_center=False,
        samplewise_center=False,
        featurewise_std_normalization=False,
        samplewise_std_normalization=False,
        rotation_range=7.0,
        width_shift_range=0.04,
        height_shift_range=0.04,
        shear_range=0.025,
        zoom_range=0.03,
        fill_mode='constant',
        cval=0,
        horizontal_flip=False,
        vertical_flip=False)
    image_datagen = ImageDataGenerator(**data_gen_args)
    mask_datagen = ImageDataGenerator(**data_gen_args)
    images = np.zeros((elements, 224, 224, 6), dtype=np.uint8)
    labels = np.zeros((elements, 224, 224, 1), dtype=np.uint8)
    image_generator = image_datagen.flow(images, batch_size=16, seed=seed)
    mask_generator = mask_datagen.flow(labels, batch_size=16, seed=seed)
    generator = combine_generator(image_generator, mask_generator)
    return images, labels, generator

max_elem = 4000
max_valid = 2000
trainImgs, trainlabels, train_generator = populate_set(elements=max_elem)
testImgs, testlabels, test_generator = populate_set(elements=max_valid)

for train1 in range(0,1000):
    # 执行训练
    # hist = model.fit(...)
    
    # 清理旧对象引用
    del trainImgs, trainlabels, train_generator
    # 强制触发垃圾回收
    gc.collect()
    
    # 加载新数据
    trainImgs, trainlabels, train_generator = populate_set(elements=max_elem)

方案3：改用tf.data.Dataset（长期推荐） #

TensorFlow的tf.data.Dataset是更现代的数据管道API，内存管理更高效，支持动态加载数据，从根本避免生成器导致的内存泄漏问题。

import tensorflow as tf
import numpy as np

def load_partial_data(elements=4000):
    # 替换为你实际加载部分数据的逻辑
    images = np.zeros((elements, 224, 224, 6), dtype=np.uint8)
    labels = np.zeros((elements, 224, 224, 1), dtype=np.uint8)
    return images, labels

def create_augmented_dataset(elements, batch_size=16):
    # 构建动态加载的数据集
    dataset = tf.data.Dataset.from_tensor_slices(load_partial_data(elements))
    
    # 实现数据增强（对应原ImageDataGenerator的配置）
    def augment(image, label):
        # 随机旋转（7度转弧度）
        image = tf.image.random_rotate(image, 7.0 / 360 * 2 * np.pi)
        # 随机平移
        image = tf.image.random_shift(image, 0.04, 0.04)
        # 随机剪切
        image = tf.keras.layers.RandomShear(0.025)(image)
        # 随机缩放
        image = tf.keras.layers.RandomZoom(0.03)(image)
        return image, label
    
    dataset = dataset.shuffle(elements).batch(batch_size).map(augment).prefetch(tf.data.AUTOTUNE)
    return dataset

max_elem = 4000
for _ in range(0,1000):
    # 每次训练前加载新数据并构建数据集
    train_dataset = create_augmented_dataset(max_elem)
    # 执行训练
    # hist = model.fit(train_dataset, epochs=epochs_together)

方案对比 #

• 方案1：最小改动现有代码，内存占用稳定，适合快速修复问题。
• 方案2：简单粗暴，但依赖Python垃圾回收机制，无法保证解决所有隐藏引用问题。
• 方案3：长期最优选择，支持分布式训练、异步加载等高级特性，需重构数据加载和增强逻辑。

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