如何对字典派生的TensorFlow Dataset做窗口化处理

问题现象 #

对字典结构派生的多特征数据集执行Dataset.window()窗口化操作时，套用单序列数据集的常规批处理写法会触发如下报错：

AttributeError: 'dict' object has no attribute 'batch'

错误触发的初始实现代码：

import tensorflow as tf
data = {'col_a': [0,1,2,3,4], 'col_b': [0,1,2,3,4]}
ds = tf.data.Dataset.from_tensor_slices(data)
windowed_ds = ds.window(size=3,shift=1, drop_remainder=True)
# 错误写法：直接对返回的字典对象调用batch方法
new_ds = windowed_ds.flat_map(lambda window: window.batch(3))

报错原因 #

字典结构的数据集经过window()方法处理后，返回的每个窗口对象是键为特征名、值为对应特征窗口数据集的字典，并非单个数据集实例，因此无法直接调用batch()方法。

预期输出 #

处理后期望得到的窗口化数据集结构如下：

{'col_a': [[0,1,2],[1,2,3],[2,3,4]], 'col_b': [[0,1,2],[1,2,3],[2,3,4]]}

修正代码 #

遍历窗口字典的每个键值对，对每个特征对应的独立窗口数据集单独执行batch()操作，再通过tf.data.Dataset.zip()将各特征的批处理结果重新打包为字典结构即可：

new_ds = windowed_ds.flat_map(
    lambda window: tf.data.Dataset.zip(
        {k: v.batch(3) for k, v in window.items()}
    )
)

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