如何适配数据增强流程避免26000+图像分割训练时出现内存溢出错误

根因说明 #

当前实现将26000+张样本全量加载到内存生成numpy数组，后续数据集拆分、ImageDataGenerator.fit操作会额外产生多份全量数据副本，叠加后超出内存上限导致notebook崩溃。另外原代码存在逻辑bug：Image.open(image).thumbnail(size)是原地修改方法，返回值为None，实际存入X_data的全是空值，即使内存足够也无法正常训练。

优化方案 #

1. 替换全量加载为流式加载（按需读盘） #

不要提前把所有图片读到numpy数组里，改用自定义生成器，训练时才逐批次读取图片、做预处理，内存中仅保留当前batch的数据，内存占用直接下降2个数量级。
示例代码如下：

import random
import numpy as np
from PIL import Image
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from sklearn.model_selection import train_test_split

# 自定义生成器，同时保证图像和掩码的增强变换一致
def batch_generator(img_paths, mask_paths, batch_size=32, img_size=(128,128), augment_args=None):
    datagen = ImageDataGenerator(**augment_args) if augment_args else ImageDataGenerator()
    while True:
        # 打乱样本顺序
        combined = list(zip(img_paths, mask_paths))
        random.shuffle(combined)
        img_paths_shuf, mask_paths_shuf = zip(*combined)
        # 逐批次加载处理
        for i in range(0, len(img_paths_shuf), batch_size):
            batch_imgs = []
            batch_masks = []
            for img_p, mask_p in zip(img_paths_shuf[i:i+batch_size], mask_paths_shuf[i:i+batch_size]):
                # 加载并处理原图
                img = Image.open(img_p).resize(img_size)
                img = np.asarray(img, dtype=np.float32) / 255.0 # 同步做归一化
                # 加载并处理分割掩码
                mask = Image.open(mask_p).resize(img_size, resample=Image.NEAREST) # 掩码用最近邻插值避免标签错误
                mask = np.asarray(mask, dtype=np.float32)
                batch_imgs.append(img)
                batch_masks.append(mask)
            batch_imgs = np.array(batch_imgs)
            batch_masks = np.array(batch_masks)
            # 同批次做增强，用相同seed保证图像和掩码变换对齐
            seed = np.random.randint(0, 10000)
            aug_imgs = datagen.flow(batch_imgs, batch_size=batch_size, seed=seed, shuffle=False).next()
            aug_masks = datagen.flow(batch_masks, batch_size=batch_size, seed=seed, shuffle=False).next()
            yield aug_imgs, aug_masks

调用方式：

# 仅读取路径，不加载图片数据
all_img_paths = sorted(glob('../input/fg_image/images/*.jpg'))
all_mask_paths = sorted(glob('../input/fg_image/masks/*.png')) # 替换为你的掩码文件路径
# 仅拆分路径，不产生数据副本
X_train_paths, X_val_paths, Y_train_paths, Y_val_paths = train_test_split(
    all_img_paths, all_mask_paths, test_size=0.2, random_state=42
)
# 定义增强参数
aug_args = dict(rotation_range=45., width_shift_range=0.1, height_shift_range=0.1)
# 初始化训练/验证生成器
train_generator = batch_generator(X_train_paths, Y_train_paths, batch_size=32, augment_args=aug_args)
val_generator = batch_generator(X_val_paths, Y_val_paths, batch_size=32, augment_args=None) # 验证集无需增强

2. 删除冗余内存操作 #

• 你当前用的几何类增强参数不需要调用ImageDataGenerator.fit()，该方法仅在做特征标准化、ZCA白化时才需要，删掉这两行可以省掉全量数据复制的内存开销。
• 不需要同时存储全量原始数据、训练集、测试集三份数组，拆分路径后按需加载即可。

调整后内存占用仅和batch_size挂钩，32的batch大小内存占用不会超过200MB，完全不会触发内存溢出。

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