如何将image_dataset_from_directory加载的数据集拆分为数据和标签？

tf.keras.preprocessing.image_dataset_from_directory返回的是tf.data.Dataset类型对象，每个迭代元素是(批次图像张量, 批次标签张量)的元组，你可以根据数据集大小选择以下两种方案拆分出独立数据和标签：

方案1：转为numpy数组（适合小数据集，可直接对接train_test_split/fit_resample） #

如果你的数据集可以全部载入内存，直接迭代整个数据集拼接所有批次即可：

import numpy as np

# 迭代所有批次，存储图像和标签
all_images = []
all_labels = []
for batch_images, batch_labels in dataset:
    # 将tf张量转为numpy数组
    all_images.append(batch_images.numpy())
    all_labels.append(batch_labels.numpy())

# 拼接所有批次得到完整的数据集
X = np.concatenate(all_images, axis=0)
y = np.concatenate(all_labels, axis=0)

操作完成后：

• X的形状为(总样本数, 578, 260, 3)，对应所有图像数据
• y的形状为(总样本数,)，对应所有样本的标签
  可以直接传入sklearn.model_selection.train_test_split或者imblearn的重采样方法使用。

方案2：tf.data原生拆分（适合大数据集，无需载入全部内存） #

如果数据集过大无法全部载入内存，不需要转为numpy数组，直接用tf.data的内置方法做拆分即可：

加载时直接拆分训练/验证集 #

在调用image_dataset_from_directory时就加入拆分参数：

train_dataset = tf.keras.preprocessing.image_dataset_from_directory(
    "train_data", 
    shuffle=True, 
    image_size=(578, 260),
    batch_size=BATCH_SIZE,
    validation_split=0.2, # 20%数据划分为验证集
    subset="training",
    seed=42 # 固定随机种子保证拆分结果可复现
)

val_dataset = tf.keras.preprocessing.image_dataset_from_directory(
    "train_data", 
    shuffle=True, 
    image_size=(578, 260),
    batch_size=BATCH_SIZE,
    validation_split=0.2,
    subset="validation",
    seed=42
)

已加载数据集的拆分 #

如果已经加载完数据集，可以用take和skip方法拆分：

# 计算总批次数
total_batches = dataset.cardinality().numpy()
train_batches = int(0.8 * total_batches) # 80%作为训练集

train_dataset = dataset.take(train_batches)
val_dataset = dataset.skip(train_batches)

注意事项 #

• 若需要使用fit_resample做类别不平衡处理，必须使用方案1转为numpy数组，imblearn的接口不支持直接传入tf.data.Dataset对象
• 大尺寸图像转numpy数组会占用较高内存，若数据集总样本数超过1万建议优先用方案2的tf.data原生操作
• 拆分前可以手动调用dataset = dataset.shuffle(1000, seed=42)打乱数据集，保证拆分均匀

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