无需拆分复制数据集，如何直接从总目录加载图像到Keras训练验证生成器

我正在开发多类别分类模型，共包含4000多个类别，对应4000个分类子文件夹，总数据集大小约30GB。当前训练模型时我需要将每个类别的图像分别复制到训练集和验证集文件夹中，以匹配分类任务的目录结构，这会额外占用30GB磁盘空间，且数据读取耗时较长。
我目前使用Keras的ImageDataGenerator API加载数据，再喂入模型训练，相关代码如下：

train_generator = train_datagen.flow_from_directory(Training_DIR, batch_size= batch_size, class_mode='categorical',target_size=(img_height,img_width))
validation_generator = validation_datagen.flow_from_directory(VALIDATION_DIR, batch_size= batch_size, class_mode='categorical',target_size=(img_height,img_width))

之后将这些生成器传入model.fit_generator函数，代码如下：

model.fit_generator(train_generator,validation_data=validation_generator)

请问是否存在更优方案，可以直接从存放所有类别子文件夹的总目录高速加载数据，无需新建目录复制图像占用双倍磁盘空间？我此前没有处理过这类大型数据集，目前复制后的文件已经占满了我的磁盘空间。
更新：我尝试了@GerryP给出的解决方案，但运行时报错如下：

第1轮/共50轮
2021-09-06 17:22:12.576079: I tensorflow/stream_executor/cuda/cuda_dnn.cc:369] 已加载cuDNN版本8101
2021-09-06 17:22:13.251294: W tensorflow/core/common_runtime/bfc_allocator.cc:272] 分配器(GPU_0_bfc)尝试分配1.19GiB显存时不足，freed_by_count=0。调用方提示这不是失败，但如果有更多显存可以提升性能。
2021-09-06 17:22:16.096112: E tensorflow/stream_executor/cuda/cuda_driver.cc:1010] 同步停止事件失败：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止
2021-09-06 17:22:16.096299: E tensorflow/stream_executor/gpu/gpu_timer.cc:55] 内部错误：销毁CUDA事件出错：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止
2021-09-06 17:22:16.097126: E tensorflow/stream_executor/gpu/gpu_timer.cc:60] 内部错误：销毁CUDA事件出错：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止
2021-09-06 17:22:16.097682: I tensorflow/stream_executor/cuda/cuda_driver.cc:732] 尝试从设备分配8B显存失败：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止
2021-09-06 17:22:16.097935: E tensorflow/stream_executor/stream.cc:4508] 内部错误：异步memset操作入队失败：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止
2021-09-06 17:22:16.098312: W tensorflow/core/kernels/gpu_utils.cc:69] 检查cudnn卷积是否存在越界读写失败，错误信息：'加载内存中CUBIN失败：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止'；跳过该检查。仅表示不会检查cudnn的越界读写，该信息仅打印一次。
2021-09-06 17:22:16.098676: I tensorflow/stream_executor/cuda/cuda_driver.cc:732] 尝试从设备分配8B显存失败：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止
2021-09-06 17:22:16.099006: E tensorflow/stream_executor/stream.cc:4508] 内部错误：异步memset操作入队失败：CUDA_ERROR_LAUNCH_TIMEOUT：启动超时，已终止
2021-09-06 17:22:16.099369: F tensorflow/stream_executor/cuda/cuda_dnn.cc:216] 检查失败：status == CUDNN_STATUS_SUCCESS (7 vs. 0) 设置cuDNN流失败。

1. 无需复制文件的数据集加载方案 #

完全不需要拆分目录复制文件，用ImageDataGenerator的flow_from_dataframe接口就可以直接读取总目录下的所有数据，自动划分训练验证集，零额外磁盘占用：
步骤：

• 遍历总数据集目录，生成包含三列信息的DataFrame：图片绝对路径、对应类别标签、划分标记（训练/验证）
• 按划分标记拆分得到训练集DataFrame和验证集DataFrame，支持分层抽样保证训练验证集的类别分布一致
• 分别调用对应的数据生成器加载两个DataFrame即可

示例代码：

import os
import pandas as pd
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 替换为你的总数据集根目录，下面是4000个类别子文件夹
ROOT_DATA_DIR = "你的总目录路径"
img_height, img_width = 224, 224
batch_size = 8 # 可根据显存情况后续调整

# 生成数据集列表
data = []
for class_name in os.listdir(ROOT_DATA_DIR):
    class_path = os.path.join(ROOT_DATA_DIR, class_name)
    if not os.path.isdir(class_path):
        continue
    for img_name in os.listdir(class_path):
        img_path = os.path.join(class_path, img_name)
        data.append({"path": img_path, "label": class_name})
df = pd.DataFrame(data)

# 随机划分80%训练，20%验证，按类别分层抽样保证分布一致
train_df, val_df = train_test_split(df, test_size=0.2, stratify=df["label"], random_state=42)

# 数据增强配置和你原来的配置保持一致即可
train_datagen = ImageDataGenerator(rescale=1./255)
val_datagen = ImageDataGenerator(rescale=1./255)

# 生成训练验证集生成器
train_generator = train_datagen.flow_from_dataframe(
    train_df,
    x_col="path",
    y_col="label",
    target_size=(img_height, img_width),
    batch_size=batch_size,
    class_mode="categorical"
)
validation_generator = val_datagen.flow_from_dataframe(
    val_df,
    x_col="path",
    y_col="label",
    target_size=(img_height, img_width),
    batch_size=batch_size,
    class_mode="categorical"
)

# TensorFlow 2.0之后已经废弃fit_generator，直接用fit即可支持生成器输入
model.fit(
    train_generator,
    validation_data=validation_generator,
    epochs=50
)

注：如果你的TensorFlow版本低于2.1，可将fit改回fit_generator，参数保持一致即可

2. CUDA报错修复方案 #

你遇到的报错本质是GPU显存不足，导致CUDA内核运行超时被系统终止，按以下顺序排查修复：

• 优先调小batch_size：先降到4甚至2，测试是否能正常启动训练，再逐步上调到显存允许的最大值
• 开启TensorFlow显存动态分配：默认TensorFlow会预占全部GPU显存，容易触发溢出，在代码最开头加入以下配置：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

• 降低输入图像分辨率：如果业务精度允许，适当缩小img_height和img_width的数值，可大幅降低显存占用
• 开启混合精度训练：启用半精度训练可以直接降低一半显存占用，同时提升训练速度，只需在代码开头加入：

from tensorflow.keras import mixed_precision
mixed_precision.set_global_policy('mixed_float16')

• 如果你是Windows平台，还可以修改WDDM的GPU超时阈值（TDR延迟），避免长耗时的内核被系统误杀。

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