如何在不增加额外维度的情况下对TF.Dataset进行批量处理？

核心原因 #

你遇到的多余维度1，是因为一开始用timeseries_dataset_from_array时设置了batch_size=1，导致每个数据集元素的特征张量自带一个长度为1的batch维度。后续再次调用batch()时，会在这个维度外层再套一层新的batch维度，最终形成(batch_size, 1, window_size, n)的形状。

两种可行解决方案 #

方案1：创建数据集时不预先设置batch_size #

直接在timeseries_dataset_from_array中把batch_size设为None（或者不指定该参数，默认就是None），这样每个数据集元素是单个样本而非批量样本，后续执行filter、map后再调用batch()，就能得到期望的形状。

修改后的关键代码：

# 创建时序数据集时不设置batch_size
Xy_ds = tf.keras.utils.timeseries_dataset_from_array(X, y, sequence_length=window_size, 
                                                     sequence_stride=1, shuffle=True)

# 后续过滤、映射操作不变
Xy_ds = Xy_ds.filter(single_source)
Xy_ds = Xy_ds.map(drop_source)

# 批量处理
Xy_ds = Xy_ds.batch(batch_size)

运行后输出的x形状会是(3, 10, 5)，符合预期。

方案2：先拆分现有批量再重新批量 #

如果因为某些原因必须保留初始的batch_size=1，可以先调用unbatch()把每个批量拆分成单个样本，再重新执行batch()：

# 创建数据集时保留batch_size=1
Xy_ds = tf.keras.utils.timeseries_dataset_from_array(X, y, sequence_length=window_size, batch_size=1, 
                                                     sequence_stride=1, shuffle=True)

# 过滤、映射操作不变
Xy_ds = Xy_ds.filter(single_source)
Xy_ds = Xy_ds.map(drop_source)

# 先拆分为单个样本，再重新批量
Xy_ds = Xy_ds.unbatch().batch(batch_size)

或者也可以在批量后用tf.squeeze去掉多余的维度：

Xy_ds = Xy_ds.batch(batch_size).map(lambda x, y: (tf.squeeze(x, axis=1), y))

验证修改后的结果 #

无论采用哪种方案，运行代码后都会得到：

batch shape:  (3, 10, 5)
total batches: 31
total datums: 93

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