在不丢失基数信息的情况下对 TensorFlow 数据集进行窗口处理？

tf.data.Dataset.window返回一个新的数据集，其元素是数据集，这些嵌套数据集的元素是所需大小的窗口。如果您有一个数据集（例如，Dataset.range(10)并想要一个像这样的窗口数据集[0 1 2] [1 2 3] ... [7 8 9]），有一个技巧可以做到这一点window plus flat_map:

>>> d = tf.data.Dataset.range(10).window(3, shift=1, drop_remainder=True).flat_map(lambda x: x.batch(3))
>>> print(list(d))
[<tf.Tensor: shape=(3,), dtype=int64, numpy=array([0, 1, 2])>, <tf.Tensor: shape=(3,), dtype=int64, numpy=array([1, 2, 3])>, <tf.Tensor: shape=(3,), dtype=int64, numpy=array([2, 3, 4])>, <tf.Tensor: shape=(3,), dtype=int64, numpy=array([3, 4, 5])>, <tf.Tensor: shape=(3,), dtype=int64, numpy=array([4, 5, 6])>, <tf.Tensor: shape=(3,), dtype=int64, numpy=array([5, 6, 7])>, <tf.Tensor: shape=(3,), dtype=int64, numpy=array([6, 7, 8])>, <tf.Tensor: shape=(3,), dtype=int64, numpy=array([7, 8, 9])>]

但是，那flat_map导致数据集丢失基数信息：

>>> d.cardinality.numpy()
<tf.Tensor: shape=(), dtype=int64, numpy=-2>

(-2 is UNKNOWN_CARDINALITY https://www.tensorflow.org/api_docs/python/tf/data#UNKNOWN_CARDINALITY; see Tensorflow 2.0：flat_map() 压平数据集的数据集返回基数 -2 https://stackoverflow.com/questions/66287320/tensorflow-2-0-flat-map-to-flatten-dataset-of-dataset-returns-cardinality-2/68281606#68281606)

我想创建此类窗口的数据集，同时保留基数信息。使用未知基数的数据集的一个小烦恼是 Keras 训练进度条需要先运行一个 epoch，然后才能生成 ETA。我试过.take(n_windows)我在哪里计算n_windows我自己，但仍然返回了一个数据集UNKNOWN_CARDINALITY.

有没有某种方法可以在不丢失基数信息的情况下对数据集进行窗口化？

主要问题是基数是静态计算的。因此无法计算 flat_map 操作的基数。您可以参考这个问题 https://github.com/tensorflow/tensorflow/issues/37998

正如您所知道的 flat_map 输入和输出的关系，解决方案是您自己设置基数。

这是有关如何设置窗口基数的示例：

import tensorflow as tf

ds = tf.data.Dataset.range(10)
print("Original cardinality -> ", ds.cardinality().numpy())
# Output:
# Original cardinality -> 10

ds = ds.window(3, shift=1, drop_remainder=True)
# cardinality at this point is still known.
# as drop_remainder is true, window cardinality will be <= original cardinality
window_cardinality = ds.cardinality()
print("window cardinality ->",window_cardinality.numpy())
# Output:
# window cardinality -> 8

ds = ds.flat_map(lambda x: x.batch(3))
# after flat_map the inferred cardinality is lost.
print("flat cardinality ->",ds.cardinality().numpy())
# Output:
# flat cardinality -> -2

# as we know the flat_map relation is 1:1 we can set the cardinality back to the original value.
ds = ds.apply(tf.data.experimental.assert_cardinality(window_cardinality))
print("dataset cardinality ->",ds.cardinality().numpy())
print("length of dataset ->", len(list(ds)))
# Output: 
# dataset cardinality -> 8
# length of dataset -> 8

for idx, x in ds.enumerate():
    print(f"{idx} -> {x}")
# Output:
# 0 -> [0 1 2]
# 1 -> [1 2 3]
# 2 -> [2 3 4]
# 3 -> [3 4 5]
# 4 -> [4 5 6]
# 5 -> [5 6 7]
# 6 -> [6 7 8]
# 7 -> [7 8 9]