使用Generator时内核重置且内存占满的问题排查及正确实现方法

首先，咱们先揪出你当前用tf.data出现内存爆、内核重置的核心原因：你设置的shuffle(buffer_size=train_ds_size)会把整个训练数据集一次性加载到内存里做打乱操作，这完全违背了流式加载节省内存的初衷！哪怕你设了batch_size，只要shuffle的buffer是全集，内存就会被直接占满，自然触发内核重置。而且测试集和验证集也没必要用训练集大小的buffer，这纯粹是额外的内存浪费。

接下来，我给你实现Keras官方推荐的Sequence类生成器——它是安全的迭代器，支持多进程，而且每次只加载一个batch的数据，完美解决内存问题：

自定义Sequence生成器实现 #

import tensorflow as tf
import numpy as np

# 加载并划分数据
cifar10_data = tf.keras.datasets.cifar10
(train_images, train_labels), (test_images, test_labels) = cifar10_data.load_data()
CLASS_NAMES= ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

validation_images, validation_labels = train_images[:5000], train_labels[:5000]
train_images, train_labels = train_images[5000:], train_labels[5000:]

# 定义预处理函数
def process_images(image, label, size=227):
    # 将图像标准化为均值为0、标准差为1
    image = tf.image.per_image_standardization(image)
    # 将图像从32x32 resize至227x227
    image = tf.image.resize(image, (227,227))
    return image, label

# 自定义Sequence生成器
class CIFAR10Sequence(tf.keras.utils.Sequence):
    def __init__(self, images, labels, batch_size=64, preprocess_func=None):
        self.images = images
        self.labels = labels
        self.batch_size = batch_size
        self.preprocess_func = preprocess_func
        # 保存数据索引，用于后续shuffle
        self.indices = np.arange(len(self.images))
    
    def __len__(self):
        # 返回总batch数（向上取整，避免遗漏最后几个样本）
        return int(np.ceil(len(self.images) / self.batch_size))
    
    def __getitem__(self, idx):
        # 获取当前batch的索引范围
        batch_indices = self.indices[idx * self.batch_size : (idx + 1) * self.batch_size]
        # 仅加载当前batch的图像和标签
        batch_images = self.images[batch_indices]
        batch_labels = self.labels[batch_indices]
        
        # 对当前batch做预处理
        if self.preprocess_func:
            processed_batch = [self.preprocess_func(img, lbl) for img, lbl in zip(batch_images, batch_labels)]
            batch_images = np.array([item[0] for item in processed_batch])
            batch_labels = np.array([item[1] for item in processed_batch])
        
        return batch_images, batch_labels
    
    def on_epoch_end(self):
        # 每个epoch结束后打乱索引，保证训练数据的随机性
        np.random.shuffle(self.indices)

# 创建生成器实例
train_generator = CIFAR10Sequence(train_images, train_labels, batch_size=64, preprocess_func=process_images)
validation_generator = CIFAR10Sequence(validation_images, validation_labels, batch_size=64, preprocess_func=process_images)
test_generator = CIFAR10Sequence(test_images, test_labels, batch_size=64, preprocess_func=process_images)

# 定义模型（和你原来的结构完全一致）
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(filters=96, kernel_size=(11,11), strides=(4,4), activation='relu', input_shape=(227,227,3)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPool2D(pool_size=(3,3), strides=(2,2)),
    tf.keras.layers.Conv2D(filters=256, kernel_size=(5,5), strides=(1,1), activation='relu', padding="same"),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPool2D(pool_size=(3,3), strides=(2,2)),
    tf.keras.layers.Conv2D(filters=384, kernel_size=(3,3), strides=(1,1), activation='relu', padding="same"),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(filters=384, kernel_size=(3,3), strides=(1,1), activation='relu', padding="same"),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(filters=256, kernel_size=(3,3), strides=(1,1), activation='relu', padding="same"),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPool2D(pool_size=(3,3), strides=(2,2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(4096, activation='relu'),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(4096, activation='relu'),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(loss='sparse_categorical_crossentropy', optimizer=tf.optimizers.SGD(learning_rate=0.001), metrics=['accuracy'])

# 用生成器训练模型（可开启多进程加速）
history = model.fit(
    train_generator,
    epochs=1,
    validation_data=validation_generator,
    verbose=1,
    validation_freq=1,
    # 可选：根据CPU核心数设置workers，开启多进程预处理
    # workers=4,
    # use_multiprocessing=True
)

关键说明 #

1. 内存友好的核心逻辑：Sequence生成器每次仅加载一个batch的数据到内存，处理完就释放，不会一次性把所有数据集读进内存，从根源上解决内存溢出问题。
2. 数据随机性保证：on_epoch_end方法会在每个epoch结束后打乱数据索引，效果和shuffle一致，但内存开销极低。
3. 多进程加速：如果你的CPU核心足够，可以打开workers和use_multiprocessing参数，让数据预处理和模型训练并行，提升整体效率。

补充：如果想继续用tf.data怎么改？ #

如果你不想换生成器，也可以修改原来的tf.data代码，把shuffle(buffer_size)改成一个合理的小值（比如1000），而不是整个数据集大小：

train_ds = (train_ds
            .map(process_images)
            .shuffle(buffer_size=1000)  # 改成1000，而非train_ds_size
            .batch(batch_size=64, drop_remainder=True))

这样tf.data只会加载1000个样本到内存做打乱，也能避免内存爆掉。

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