如何将含输入与标签的tf.data.Dataset转换为顺序一致的NumPy数组

当前使用TensorFlow 2.9.1版本，持有类型为tf.data.Dataset的test_dataset对象，数据集同时存储输入数据与标签：输入为4维张量，标签为3维张量，数据集结构如下：

print(test_dataset)
<PrefetchDataset element_spec=(TensorSpec(shape=(64, 5, 548, 1), dtype=tf.float64, name=None), TensorSpec(shape=(64, 1, 1), dtype=tf.float64, name=None))>

张量第一维度为小批量（minibatch）大小。需求是将该TensorFlow数据集转换为两个NumPy数组：存储输入的X_test和存储标签的y_test，二者样本顺序必须完全一致——即(X_test[0], y_test[0])对应test_dataset中的第一个样本，且沿张量第一维度（批次维度）拼接所有批次结果。

方法1：直接使用np.concatenate分两次迭代 #

原有实现代码：

X_test = np.concatenate([x for x, _ in test_dataset], axis=0)
y_test = np.concatenate([y for _, y in test_dataset], axis=0)

该方案存在两个明确缺陷：

• 需要对同一个数据集迭代两次，存在不必要的算力浪费
• X_test与y_test的样本顺序无法对齐。实际测试验证：连续两次运行X_test = np.concatenate([x for x, _ in test_dataset], axis=0)和X_test2 = np.concatenate([x for x, _ in test_dataset], axis=0)得到的两个数组形状一致但内容不同，原因是数据集迭代过程如果带有shuffle逻辑，两次独立迭代得到的样本顺序不匹配，最终生成的输入和标签数组无法对应。

方法2：使用tfds.as_numpy转换 #

根据接口说明，tfds.as_numpy可将TensorFlow数据集转换为NumPy数组的可迭代对象，基础调用代码：

import tensorflow_datasets as tfds
np_test_dataset = tfds.as_numpy(test_dataset)
print(np_test_dataset)
<generator object _eager_dataset_iterator at 0x7fee81fd8b30>

该方法的问题是不清楚后续如何处理返回的生成器对象，将其转换为形状正确、顺序对齐的两个独立NumPy数组。

核心思路是单次迭代数据集，同时收集输入和标签批次，最后统一做拼接，从根源避免两次迭代带来的顺序错位和算力浪费。

方案1：原生TensorFlow+NumPy实现（推荐，无额外依赖） #

单次遍历数据集，分别存储每个批次的输入和标签，最后沿批次维度拼接：

import numpy as np

x_batches = []
y_batches = []

# 单次迭代同时取输入和标签，从流程上保证顺序完全对齐
for x_batch, y_batch in test_dataset:
    x_batches.append(x_batch.numpy())
    y_batches.append(y_batch.numpy())

# 沿第0维（批次维度）拼接所有批次
X_test = np.concatenate(x_batches, axis=0)
y_test = np.concatenate(y_batches, axis=0)

注意：如果你的测试数据集pipeline本身带有shuffle操作，请在转换前移除测试集的shuffle逻辑，或者给shuffle传入固定随机种子，否则即使单次迭代，每次运行数据集的输出顺序也会随机变化。

方案2：基于tfds.as_numpy的实现 #

如果你已经使用tfds.as_numpy将数据集转为NumPy生成器，只需要单次遍历该生成器收集批次再拼接即可，逻辑和方案1完全一致：

import numpy as np
import tensorflow_datasets as tfds

np_test_dataset = tfds.as_numpy(test_dataset)
x_batches = []
y_batches = []

for x_batch, y_batch in np_test_dataset:
    # 此处x_batch、y_batch本身已经是NumPy数组，不需要再调用.numpy()转换
    x_batches.append(x_batch)
    y_batches.append(y_batch)

X_test = np.concatenate(x_batches, axis=0)
y_test = np.concatenate(y_batches, axis=0)

结果校验 #

转换完成后可以做简单校验确认正确性：

# 校验形状：输入形状应为(总样本数, 5, 548, 1)，标签形状应为(总样本数, 1, 1)
print(X_test.shape, y_test.shape)

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