上传图像数据集后Google Colab内存占满运行崩溃如何解决？

手动完成全量数据预处理时，11000张图像的预处理结果会全部常驻内存，直接占满Colab的运行内存配额。Keras的flow_from_directory默认采用流式加载逻辑，每次仅读取当前训练所需的一个batch的图像，预处理完成喂入模型后就会释放对应内存，因此不会出现内存溢出问题。

• 优先复用flow_from_directory能力，无需自行实现数据加载逻辑。你可以通过接口的preprocessing_function参数传入自定义预处理函数，既可以保留原有的预处理逻辑，也能享受到原生的内存优化特性，改造成本最低。
• 如果需要自定义数据加载流程，不要全量加载所有图像，改用生成器实现分批加载，每次迭代仅处理一个batch的数据，处理完成的批次不保留额外引用，交由系统自动回收内存，参考实现如下：

import os
import cv2
import random
import numpy as np
import tensorflow as tf
import gc

def custom_data_generator(img_root_dir, batch_size, img_size=(224,224), num_classes=11):
    class_folders = sorted(os.listdir(img_root_dir))
    while True:
        # 每个epoch重新打乱样本顺序
        all_img_paths = []
        all_labels = []
        for cls_idx, cls_folder in enumerate(class_folders):
            cls_folder_path = os.path.join(img_root_dir, cls_folder)
            cls_img_paths = [os.path.join(cls_folder_path, img_name) for img_name in os.listdir(cls_folder_path)]
            all_img_paths.extend(cls_img_paths)
            all_labels.extend([cls_idx] * len(cls_img_paths))
        # 打乱样本
        combined = list(zip(all_img_paths, all_labels))
        random.shuffle(combined)
        all_img_paths, all_labels = zip(*combined)
        # 分批加载处理
        for step in range(0, len(all_img_paths), batch_size):
            batch_paths = all_img_paths[step:step+batch_size]
            batch_labels = all_labels[step:step+batch_size]
            batch_imgs = []
            for img_path in batch_paths:
                img = cv2.imread(img_path)
                img = cv2.resize(img, img_size)
                # 此处替换为你自己的预处理逻辑
                img = img / 255.0
                batch_imgs.append(img)
            yield np.array(batch_imgs), tf.keras.utils.to_categorical(batch_labels, num_classes=num_classes)

• 额外优化技巧：处理过程中不需要的中间变量及时用del删除，再调用gc.collect()手动触发垃圾回收，减少无效内存占用。如果从Google Drive加载数据耗时久，可以先把数据集打包为zip压缩包上传到Drive，在Colab中先将压缩包复制到本地磁盘再解压，相比直接读取Drive内的零散文件速度提升非常明显。

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