如何使用tf.keras.utils.image_dataset_from_directory划分80%训练、10%验证、10%测试集？

可以通过两次拆分数据集来达到目标比例，下面给出两种可行方案：

方案一：分两次调用image_dataset_from_directory拆分 #

第一步：先拆分出测试集（10%）和训练+验证池（90%） #

从全量数据中先剥离出10%作为测试集，剩下90%作为后续拆分训练、验证集的基础：

# 加载训练+验证池（占总数据的90%）
train_val_dataset = tf.keras.utils.image_dataset_from_directory(
    directory,
    shuffle=True,
    batch_size=BATCH_SIZE,
    image_size=IMG_SIZE,
    validation_split=0.1,
    subset='training',
    seed=42
)

# 加载测试集（占总数据的10%）
test_dataset = tf.keras.utils.image_dataset_from_directory(
    directory,
    shuffle=True,
    batch_size=BATCH_SIZE,
    image_size=IMG_SIZE,
    validation_split=0.1,
    subset='validation',
    seed=42
)

第二步：从训练+验证池中拆分训练集（80%总数据）和验证集（10%总数据） #

训练+验证池占总数据的90%，我们需要从中拆分出1/9作为验证集（对应总数据的10%），剩余8/9作为训练集（对应总数据的80%）：

# 计算训练+验证池的总样本数
total_train_val = len(train_val_dataset) * BATCH_SIZE
# 计算验证集需要的样本量
val_sample_count = total_train_val // 9

# 拆分验证集
validation_dataset = train_val_dataset.take(val_sample_count)
# 拆分训练集（跳过验证集的样本）
train_dataset = train_val_dataset.skip(val_sample_count)

方案二：加载全量数据集后直接拆分 #

如果不想重复调用加载函数，也可以先加载完整数据集，再按比例拆分：

# 加载全量数据集
full_dataset = tf.keras.utils.image_dataset_from_directory(
    directory,
    shuffle=True,
    batch_size=BATCH_SIZE,
    image_size=IMG_SIZE,
    seed=42
)

# 计算各数据集的样本数量
total_samples = len(full_dataset) * BATCH_SIZE
train_size = int(0.8 * total_samples)
val_size = int(0.1 * total_samples)

# 依次拆分训练、验证、测试集
train_dataset = full_dataset.take(train_size)
remaining_data = full_dataset.skip(train_size)
validation_dataset = remaining_data.take(val_size)
test_dataset = remaining_data.skip(val_size)

注意事项 #

• 所有涉及数据打乱的操作要保持相同的seed，确保划分的一致性。
• 如果数据集样本数无法被batch_size整除，可以用full_dataset.cardinality().numpy()获取准确的批次数量，再结合batch_size计算样本数，避免误差。

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