TensorFlow
Padded Batches
tf.train.SequenceExample
DataSet API
Machine Learning

How do I create padded batches in Tensorflow for tf.train.SequenceExample data using the DataSet API?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

tf.train.SequenceExample is a common format for variable-length sequence data such as token sequences, frame features, or event streams. When those sequences have different lengths, the normal batch operation is not enough, so the standard TensorFlow solution is to parse each example first and then use padded_batch to align sequence lengths within each batch.

Build or Read Serialized SequenceExample Records

A SequenceExample usually contains context features plus one or more variable-length feature lists. The example below creates two serialized records so the padding behavior is easy to see.

python
1import tensorflow as tf
2
3
4def make_example(label, values):
5    ex = tf.train.SequenceExample()
6    ex.context.feature["label"].int64_list.value.append(label)
7    fl = ex.feature_lists.feature_list["tokens"]
8    for value in values:
9        fl.feature.add().int64_list.value.append(value)
10    return ex.SerializeToString()
11
12
13serialized = [
14    make_example(0, [10, 11, 12]),
15    make_example(1, [20, 21]),
16]
17
18dataset = tf.data.Dataset.from_tensor_slices(serialized)

In a real pipeline these records often come from TFRecord files, but the padded-batch logic is the same.

Parse the Sequence Data First

Use tf.io.parse_single_sequence_example inside Dataset.map so the dataset contains tensors rather than raw serialized bytes.

python
1context_features = {
2    "label": tf.io.FixedLenFeature([], tf.int64),
3}
4sequence_features = {
5    "tokens": tf.io.FixedLenSequenceFeature([], tf.int64),
6}
7
8
9def parse_record(example_bytes):
10    context, sequence = tf.io.parse_single_sequence_example(
11        example_bytes,
12        context_features=context_features,
13        sequence_features=sequence_features,
14    )
15    return {
16        "label": context["label"],
17        "tokens": sequence["tokens"],
18    }
19
20
21parsed = dataset.map(parse_record)

After parsing, each element has a scalar label and a one-dimensional tokens tensor whose length may differ between records.

Use padded_batch

Now batch the parsed dataset and tell TensorFlow which shapes may vary and what padding value to use.

python
1batched = parsed.padded_batch(
2    batch_size=2,
3    padded_shapes={
4        "label": [],
5        "tokens": [None],
6    },
7    padding_values={
8        "label": tf.constant(0, dtype=tf.int64),
9        "tokens": tf.constant(0, dtype=tf.int64),
10    },
11)
12
13for batch in batched:
14    print(batch)

[None] means the sequence length may vary and should be padded to the longest sequence in the batch. The scalar label uses [] because it does not need sequence padding.

What the Shapes Mean

For the two-example dataset above, the result is a batch where:

  • 'label has shape (2,)'
  • 'tokens has shape (2, 3)'

The second example had only two tokens, so TensorFlow pads it with the provided value 0 to match the longest sequence length in the batch.

That is the main idea behind padded batching: keep batch dimensions regular without requiring all examples to have identical raw sequence lengths.

Include Additional Sequence Features

If your SequenceExample contains several parallel sequences, add them all to sequence_features and list them in padded_shapes.

python
1sequence_features = {
2    "tokens": tf.io.FixedLenSequenceFeature([], tf.int64),
3    "weights": tf.io.FixedLenSequenceFeature([], tf.float32),
4}
5
6padded_shapes = {
7    "label": [],
8    "tokens": [None],
9    "weights": [None],
10}

The important rule is that each variable-length dimension needs a None in the padded shape declaration.

Pipeline Tips

A typical input pipeline also adds performance helpers.

python
1batched = (
2    dataset
3    .map(parse_record, num_parallel_calls=tf.data.AUTOTUNE)
4    .padded_batch(32, padded_shapes={"label": [], "tokens": [None]})
5    .prefetch(tf.data.AUTOTUNE)
6)

prefetch and parallel mapping do not change correctness, but they often improve training throughput.

Common Pitfalls

The most common mistake is calling batch instead of padded_batch on variable-length sequences. Regular batching requires equal shapes and will fail.

Another issue is using the wrong padded_shapes declaration. Scalars need [], while variable-length one-dimensional sequences need [None].

A third problem is forgetting to parse the serialized SequenceExample before batching. padded_batch works on tensors, not on raw serialized protocol-buffer bytes.

Summary

  • Parse SequenceExample records with tf.io.parse_single_sequence_example before batching.
  • Use padded_batch when sequence lengths vary across examples.
  • Mark variable-length dimensions with None in padded_shapes.
  • Supply explicit padding_values when the default zero is not appropriate.
  • Add prefetch and parallel mapping after the logic is correct.

Course illustration
Course illustration

All Rights Reserved.