TensorFlow中Dataset.batch处理zip数据集不符合预期的问题及疑问

问题背景 #

创建如下TensorFlow数据集：

import tensorflow as tf

a = tf.data.Dataset.range(1, 16)
b = tf.data.Dataset.range(16, 32)
zipped = tf.data.Dataset.zip((a, b))
list(zipped.as_numpy_iterator())

# 输出: 
[(0, 16),
 (1, 17),
 (2, 18),
 (3, 19),
 (4, 20),
 (5, 21),
 (6, 22),
 (7, 23),
 (8, 24),
 (9, 25),
 (10, 26),
 (11, 27),
 (12, 28),
 (13, 29),
 (14, 30),
 (15, 31)]

对其应用batch(4)时，预期每个批次是包含4个元组的数组：

[[(0, 16), (1, 17), (2, 18), (3, 19)],
 [(4, 20), (5, 21), (6, 22), (7, 23)],
 [(8, 24), (9, 25), (10, 26), (11, 27)],
 [(12, 28), (13, 29), (14, 30), (15, 31)]]

但实际得到的结果是：

batched = zipped.batch(4)
list(batched.as_numpy_iterator())

# 输出:
[(array([0, 1, 2, 3]), array([16, 17, 18, 19])), 
 (array([4, 5, 6, 7]), array([20, 21, 22, 23])), 
 (array([ 8,  9, 10, 11]), array([24, 25, 26, 27])), 
 (array([12, 13, 14, 15]), array([28, 29, 30, 31]))]

查阅文档得知这是预期行为：

The components of the resulting element will have an additional outer dimension, which will be batch_size

解答 #

一、Dataset.batch 设计目的 #

这种实现完全适配TensorFlow的深度学习计算范式：

• 模型输入通常要求特征与标签分离为独立张量，且每个张量的第一维度为batch size。比如你的场景中，模型需要接收形状为(batch_size,)的特征张量和标签张量，而非元组数组，这样能直接对接后续的模型层、损失函数，无需额外拆分操作，避免性能损耗。
• 保持数据集元素结构一致性：原数据集元素是(特征, 标签)元组，批处理后依然是元组，仅内部张量新增batch维度，符合TensorFlow对数据输入的标准化要求。

二、实现预期结果的替代方法 #

如果需要得到“批次内是元组数组”的结构，可通过以下两种方式实现：

方法1：batch后重组张量

先执行默认批处理，再将每个批次的两个张量重组为元组数组：

batched = zipped.batch(4)
# 将批次内的特征、标签张量堆叠为二维数组，再转换为元组列表
result = [list(map(tuple, tf.stack((x, y), axis=1).numpy())) for x, y in batched.as_numpy_iterator()]
print(result)

输出与预期一致：

[[(0, 16), (1, 17), (2, 18), (3, 19)],
 [(4, 20), (5, 21), (6, 22), (7, 23)],
 [(8, 24), (9, 25), (10, 26), (11, 27)],
 [(12, 28), (13, 29), (14, 30), (15, 31)]]

方法2：先打包元素再batch

先将每个元组元素打包为单一张量，再执行批处理，最后转换为元组列表：

# 将每个(特征,标签)元组打包成单个张量
packed = zipped.map(lambda x, y: tf.stack([x, y]))
batched_packed = packed.batch(4)
# 转换为预期的元组数组结构
result = [list(map(tuple, batch.numpy())) for batch in batched_packed.as_numpy_iterator()]
print(result)

同样能得到预期结果。

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