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

user_133374 0
Viewed 463

E.g. imagine I use the Librispeech dataset via TFDS (or whatever dataset, including sequences of varying length of data), and then use padded_batch to create batches, e.g. like this:

import tensorflow_datasets as tfds

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

Now when iterating through the resulting dataset, i.e. over the (padded) batches, how would I know the original sequence lengths in the padded batch? Or is this information lost at this point? How would I extend the pipeline to include it? Is there a special dataset like AddSeqLengthInfoDataset or so? This would need to run before the padded_batch, right? (This is basically an equivalent of my question for TF PaddingFIFOQueue but for tf.data.Dataset.)

Is there some example? (I wonder a bit that I have not found anything about this. I would assume this is a pretty standard requirement when you work on sequences, that you need to have the information about the original sequence lengths, or not?)

tensorflow-datasets tensorflow
1 Answers

You can just add a new field to the dataset holding the size of the sequence, for example like this:

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)

Output:

{'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]}

Then you can use for example tf.sequence_mask with the value of that field to mask the padding values.

Another option is simply to pass some special padding_values to padded_batch that cannot appear in the actual data, e.g. -1 or nan, but that depends on whether those are actually invalid values for your problem.

user_17827920
Related