Introduction

tf.data.Dataset.from_generator is useful when data loading logic is easier in Python than in pure TensorFlow ops. Variable batch size is common when last batch is smaller or when upstream producer emits adaptive chunks. The key to making this stable is a correct output_signature with flexible batch axis and strict consistency for all other dimensions.

Define a Precise Generator Contract

Before creating the dataset, define exactly what each yield returns:

• Structure, such as tuple of features and labels.
• Dtype for each tensor.
• Which dimensions may vary.

Example generator that emits different batch sizes:

1import numpy as np
2import tensorflow as tf
3
4
5def batch_gen():
6    for b in [4, 2, 5]:
7        x = np.random.randn(b, 10).astype("float32")
8        y = np.random.randint(0, 2, size=(b, 1)).astype("float32")
9        yield x, y

If you skip this contract step, shape and dtype errors appear later and are harder to trace.

Use None for Variable Leading Axis

Declare flexible batch dimension in TensorSpec.

1output_signature = (
2    tf.TensorSpec(shape=(None, 10), dtype=tf.float32),
3    tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
4)
5
6ds = tf.data.Dataset.from_generator(batch_gen, output_signature=output_signature)
7
8for xb, yb in ds:
9    print(xb.shape, yb.shape)

Only the leading axis should usually be flexible. Feature dimension and label dimension should remain fixed unless your model supports variable-length inputs explicitly.

Decide Where Batching Happens

Two valid patterns exist:

• Generator yields full batches.
• Generator yields single examples and dataset does .batch.

Do not mix both accidentally. Double-batching is a common source of rank errors.

Single-example generator pattern:

1def example_gen():
2    for _ in range(12):
3        x = np.random.randn(10).astype("float32")
4        y = np.random.randint(0, 2, size=(1,)).astype("float32")
5        yield x, y
6
7sig = (
8    tf.TensorSpec(shape=(10,), dtype=tf.float32),
9    tf.TensorSpec(shape=(1,), dtype=tf.float32),
10)
11
12ds = tf.data.Dataset.from_generator(example_gen, output_signature=sig)
13ds = ds.batch(4, drop_remainder=False)

Keep one batching stage for predictable shapes.

Variable Batch Size vs Variable Sequence Length

These are different concerns:

• Variable batch size means first axis changes.
• Variable sequence length means inner time dimension changes.

If sequences vary in length, use padding or ragged handling.

1def seq_gen():
2    yield tf.constant([1, 2, 3], dtype=tf.int32)
3    yield tf.constant([4, 5], dtype=tf.int32)
4
5seq_ds = tf.data.Dataset.from_generator(
6    seq_gen,
7    output_signature=tf.TensorSpec(shape=(None,), dtype=tf.int32)
8)
9
10padded = seq_ds.padded_batch(2, padded_shapes=[None], padding_values=0)
11for batch in padded:
12    print(batch)

Treating variable sequence length as variable batch size leads to incorrect model input handling.

Improve Throughput and Stability

Python generators run on host side, so heavy logic can bottleneck training. Keep generator minimal and move transform-heavy steps to tf.data ops when possible.

Add prefetch for pipeline overlap:

ds = ds.prefetch(tf.data.AUTOTUNE)

Also cast values inside generator to avoid dtype drift from mixed sources.

x = np.asarray(x, dtype=np.float32)

Validation Before Training

Before feeding into model:

• Iterate through a few batches.
• Print shapes and dtypes.
• Confirm last partial batch behavior.
• Validate label alignment.

Quick check:

for xb, yb in ds.take(2):
    print(xb.dtype, xb.shape, yb.dtype, yb.shape)

This catches most contract mismatches early.

Common Pitfalls

• Declaring fixed batch size in signature while yielding variable sizes. Fix by using None for leading axis.
• Applying .batch to data already batched by generator. Fix by choosing one batching strategy.
• Mixing variable sequence length without padding. Fix with padded_batch or ragged tensors.
• Yielding inconsistent dtypes across iterations. Fix by explicit casting in generator.
• Putting expensive preprocessing inside Python generator. Fix by moving transforms into tf.data pipeline ops.

Summary

• Variable batch size works with from_generator when output_signature is defined correctly.
• Keep non-batch dimensions and dtypes consistent.
• Separate batching responsibilities clearly to avoid rank bugs.
• Handle variable sequence length with padding strategies, not ad hoc shape changes.
• Validate shapes early and tune pipeline throughput with prefetch.