Python tf.data.TFRecordDataset.window用法及代码示例

用法

window(
    size, shift=None, stride=1, drop_remainder=False, name=None
)

返回"windows" 的数据集。

每个"window" 都是一个数据集，其中包含输入数据集元素的子集。这些是大小为 size 的有限数据集(如果没有足够的输入元素来填充窗口并且 drop_remainder 评估为 False ，则可能更少)。

例如：

dataset = tf.data.Dataset.range(7).window(3)
for window in dataset:
  print(window)
<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int64, name=None)>
<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int64, name=None)>
<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int64, name=None)>

由于窗口是数据集，因此可以对其进行迭代：

for window in dataset:
  print([item.numpy() for item in window])
[0, 1, 2]
[3, 4, 5]
[6]

转移

shift 参数确定要在每个窗口的开始之间移动的输入元素的数量。如果窗口和元素都从 0 开始编号，则窗口 k 中的第一个元素将是输入数据集的元素 k * shift。特别是，第一个窗口的第一个元素将始终是输入数据集的第一个元素。

dataset = tf.data.Dataset.range(7).window(3, shift=1,
                                          drop_remainder=True)
for window in dataset:
  print(list(window.as_numpy_iterator()))
[0, 1, 2]
[1, 2, 3]
[2, 3, 4]
[3, 4, 5]
[4, 5, 6]

跨步

stride 参数确定窗口内输入元素之间的步幅。

dataset = tf.data.Dataset.range(7).window(3, shift=1, stride=2,
                                          drop_remainder=True)
for window in dataset:
  print(list(window.as_numpy_iterator()))
[0, 2, 4]
[1, 3, 5]
[2, 4, 6]

嵌套元素

当window 转换应用于元素是嵌套结构的数据集时，它会生成一个数据集，其中元素具有相同的嵌套结构，但每个叶子都被一个窗口替换。换句话说，嵌套应用在窗口之外，而不是在窗口内部。

类型签名是：

def window(
    self:Dataset[Nest[T]], ...
) -> Dataset[Nest[Dataset[T]]]

将 window 应用于 Dataset 的元组给出一个窗口元组：

dataset = tf.data.Dataset.from_tensor_slices(([1, 2, 3, 4, 5],
                                              [6, 7, 8, 9, 10]))
dataset = dataset.window(2)
windows = next(iter(dataset))
windows
(<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int32, name=None)>,
 <...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int32, name=None)>)

def to_numpy(ds):
  return list(ds.as_numpy_iterator())

for windows in dataset:
  print(to_numpy(windows[0]), to_numpy(windows[1]))
[1, 2] [6, 7]
[3, 4] [8, 9]
[5] [10]

将 window 应用于字典的 Dataset 会给出 Datasets 的字典：

dataset = tf.data.Dataset.from_tensor_slices({'a':[1, 2, 3],
                                              'b':[4, 5, 6],
                                              'c':[7, 8, 9]})
dataset = dataset.window(2)
def to_numpy(ds):
  return list(ds.as_numpy_iterator())

for windows in dataset:
  print(tf.nest.map_structure(to_numpy, windows))
{'a':[1, 2], 'b':[4, 5], 'c':[7, 8]}
{'a':[3], 'b':[6], 'c':[9]}

展平窗口数据集

Dataset.flat_map 和 Dataset.interleave 方法可用于将窗口数据集展平为单个数据集。

flat_map 的参数是一个从数据集中获取元素并返回 Dataset 的函数。 flat_map 按顺序将生成的数据集链接在一起。

例如，要将每个窗口变成密集张量：

size = 3
dataset = tf.data.Dataset.range(7).window(size, shift=1,
                                          drop_remainder=True)
batched = dataset.flat_map(lambda x:x.batch(3))
for batch in batched:
  print(batch.numpy())
[0 1 2]
[1 2 3]
[2 3 4]
[3 4 5]
[4 5 6]