如何从image_dataset_from_directory生成的train_data提取x_train与y_train

问题描述 #

现有通过如下代码生成的train_data数据集对象：

train_data = tf.keras.preprocessing.image_dataset_from_directory(
    directory = train_dir,
    image_size = IMG_SIZE,
    label_mode = 'categorical',
    batch_size = BATCH_SIZE,
    shuffle = True,
    subset = 'training',
    validation_split = 0.15,
    seed=42)

请问如何从中提取出x_train（训练特征数据）与y_train（训练标签数据）？

解决方法 #

方法一：迭代拼接批次（内存友好） #

train_data是按批次存储的tf.data.Dataset对象，你可以遍历所有批次，将特征和标签分别拼接成完整张量：

import tensorflow as tf

# 初始化空列表存储每一批的特征和标签
x_train_list = []
y_train_list = []

# 遍历数据集的所有批次
for x_batch, y_batch in train_data:
    x_train_list.append(x_batch)
    y_train_list.append(y_batch)

# 拼接所有批次得到完整的特征与标签张量
x_train = tf.concat(x_train_list, axis=0)
y_train = tf.concat(y_train_list, axis=0)

方法二：直接转为单批次（内存足够时用） #

如果数据集大小在内存承载范围内，可以直接将数据集合并为单个批次，再提取特征和标签：

# 将数据集转为单个批次（需确保内存能装下所有数据）
single_batch_dataset = train_data.unbatch().batch(train_data.cardinality().numpy())
x_train, y_train = next(iter(single_batch_dataset))

# 可选：转换为NumPy数组（根据后续需求选择）
x_train = x_train.numpy()
y_train = y_train.numpy()

注意事项 #

• 因为你设置了shuffle=True，提取出的x_train和y_train会是打乱后的顺序，和原数据集的打乱逻辑一致。
• 若数据集规模较大，优先用方法一，避免内存溢出问题。

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