如何用ImageDataGenerator从训练集中按类别随机抽取10%数据子集？

实现方式 #

直接通过ImageDataGenerator自带的validation_split参数即可实现需求，不需要修改现有目录结构，也不需要额外写采样逻辑。
该参数会自动按类别等比例分层抽样，完美匹配你每个类别抽取10%样本的要求。

代码示例 #

from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_data_gen = ImageDataGenerator(
    rescale=1./255,
    validation_split=0.9 # 预留90%数据为验证集，剩下的10%为可用训练子集
)

train_data_10_percent_subset = train_data_gen.flow_from_directory(
    directory='/root_dir/train/',
    target_size=(224, 224),
    class_mode='categorical',
    batch_size=32,
    seed=42, # 固定seed保证抽样结果可复现
    subset='training' # 取拆分后的训练子集，即总训练数据的10%
)

# 测试集加载逻辑保持不变
test_data_gen = ImageDataGenerator(rescale=1./255)
test_data = test_data_gen.flow_from_directory(
    directory='/root_dir/test/',
    target_size=(224, 224),
    class_mode='categorical',
    batch_size=32,
    seed=42
)

效果验证 #

你可以打印train_data_10_percent_subset.samples查看总样本量，正常返回值为225，对应3个类别每个类别75张样本，完全符合需求。

可选写法 #

如果觉得上述参数逻辑不好理解，也可以换如下写法，效果完全一致：

train_data_gen = ImageDataGenerator(
    rescale=1./255,
    validation_split=0.1 # 预留10%数据为验证集
)

train_data_10_percent_subset = train_data_gen.flow_from_directory(
    directory='/root_dir/train/',
    target_size=(224, 224),
    class_mode='categorical',
    batch_size=32,
    seed=42,
    subset='validation' # 直接取拆分后的验证集，即总训练数据的10%
)

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