如何从 `tf.data.Dataset` 的填充批次中获取序列长度？

Question

例如假设我使用 Librispeech dataset via TFDS (or whatever dataset, including sequences of varying length of data), and then use padded_batch 来创建批次，例如像这样：

import tensorflow_datasets as tfds

dataset = tfds.load(name="librispeech", split="train_clean100")
dataset = dataset.shuffle(1024)
dataset = dataset.padded_batch(32)

现在，当遍历生成的数据集时，即在（填充的）批次上，我如何知道填充批次中的原始序列长度？还是此时此信息丢失了？我将如何扩展管道以包含它？是否有像 AddSeqLengthInfoDataset 这样的特殊数据集？这样就需要运行beforepadded_batch了吧？（这基本上等同于，但 tf.data.Dataset。）

有例子吗？（我有点奇怪，我还没有找到任何关于这个的信息。我认为这是一个非常标准的要求，当你处理序列时，你需要有关于原始序列长度的信息，还是不需要？）

Answer 1

您可以向数据集中添加一个新字段来保存序列的大小，例如：

import tensorflow as tf

# Make a dataset with variable-size data
def generate_data():
    for i in range(10):
        yield {'id': i, 'data': range(i % 5)}
ds = tf.data.Dataset.from_generator(generate_data,
                                    {'id': tf.int32, 'data': tf.int32},
                                    {'id': [], 'data': [None]})
# Add field with size of data
ds = ds.map(lambda item: {**item, 'size': tf.shape(item['data'])[0]})
# Padded batch
ds = ds.padded_batch(3)
# Show dataset
for batch in ds:
    tf.print(batch)

输出：

{'data': [[0 0]
 [0 0]
 [0 1]], 'id': [0 1 2], 'size': [0 1 2]}
{'data': [[0 1 2 0]
 [0 1 2 3]
 [0 0 0 0]], 'id': [3 4 5], 'size': [3 4 0]}
{'data': [[0 0 0]
 [0 1 0]
 [0 1 2]], 'id': [6 7 8], 'size': [1 2 3]}
{'data': [[0 1 2 3]], 'id': [9], 'size': [4]}

然后您可以使用 tf.sequence_mask 和该字段的值来屏蔽填充值。

另一种选择是将一些特殊的 padding_values 传递给 padded_batch，这些特殊的 padding_values 不会出现在实际数据中，例如-1 或 nan，但这取决于它们是否实际上是您问题的无效值。

如何从 `tf.data.Dataset` 的填充批次中获取序列长度？

How to get the sequence lengths from padded batches from a `tf.data.Dataset`?

tensorflow

tensorflow-datasets