如何从TensorFlow prefetch tf.data.Dataset提取X、Y为numpy数组

核心逻辑是遍历数据集时分批收集特征和标签，最后拼接为完整数组，输出格式和train_test_split()返回的X_train/Y_train完全一致，可直接传入只接收numpy数组的函数。

具体实现步骤 #

• 提前导入numpy依赖
• 初始化两个空列表，分别暂存每个批次的X特征、Y标签（列表逐批追加的效率远高于直接拼接大数组）
• 遍历prefetch数据集，逐批取出X、Y，转为numpy格式后追加到对应列表
• 沿样本维度拼接所有批次的数组，得到最终的完整numpy变量

可直接运行的代码示例 #

import numpy as np

# 初始化暂存列表
X_list = []
Y_list = []

# 遍历prefetch数据集
for item in train_dataset:
    X_batch, Y_batch = item[0], item[1]
    # 如果批次数据是张量格式（TensorFlow/PyTorch张量都适用），直接转numpy即可
    # 若为PyTorch GPU张量，需要先移到CPU再转换：X_batch.cpu().numpy()
    X_list.append(np.array(X_batch))
    Y_list.append(np.array(Y_batch))

# 拼接为完整数组，和sklearn输出格式完全对齐
X_train = np.concatenate(X_list, axis=0)
Y_train = np.concatenate(Y_list, axis=0)

结果校验 #

执行完代码后可以打印数组形状确认结果正确：

print(X_train.shape)  # 输出格式应为(总样本数, 特征维度...)
print(Y_train.shape)  # 输出格式应为(总样本数, )或(总样本数, 标签维度...)

注意：如果你的prefetch数据集开启了多epoch重复配置，遍历前记得先关闭重复设置，否则会重复采集样本导致数组长度超出预期。

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