Introduction

Stratified sampling means preserving class balance while you split or sample data. In TensorFlow, that requirement shows up in two slightly different forms: you may want a stratified train-validation split, or you may want balanced samples during training. tf.data can support both, but the implementation depends on which of those goals you actually have.

Use a Stratified Split Before Building the Dataset

If your main goal is to create training and validation sets with similar class ratios, the easiest approach is usually to stratify before you turn the data into a tf.data.Dataset.

1import tensorflow as tf
2from sklearn.model_selection import train_test_split
3
4features = tf.constant([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]])
5labels = tf.constant([0, 0, 0, 0, 1, 1])
6
7x_train, x_val, y_train, y_val = train_test_split(
8    features.numpy(),
9    labels.numpy(),
10    test_size=0.33,
11    stratify=labels.numpy(),
12    random_state=42,
13)
14
15train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(2)
16val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(2)

This is often the cleanest answer. The split logic stays simple, and the TensorFlow pipeline only handles batching, mapping, and prefetching.

Build Balanced Sampling Pipelines with tf.data

If you want class-balanced training batches from an imbalanced source, create one dataset per class and combine them with sample_from_datasets.

1import tensorflow as tf
2
3features = tf.constant([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]])
4labels = tf.constant([0, 0, 0, 0, 1, 1])
5
6base_ds = tf.data.Dataset.from_tensor_slices((features, labels))
7
8class_0 = base_ds.filter(lambda x, y: tf.equal(y, 0))
9class_1 = base_ds.filter(lambda x, y: tf.equal(y, 1))
10
11balanced_ds = tf.data.Dataset.sample_from_datasets(
12    [class_0.repeat(), class_1.repeat()],
13    weights=[0.5, 0.5]
14).batch(2)
15
16for batch_features, batch_labels in balanced_ds.take(3):
17    print(batch_labels.numpy())

This does not preserve the original distribution. Instead, it creates a balanced stream, which is often what you want for training on imbalanced labels.

Know the Difference Between Stratified and Balanced

These terms are related but not identical:

• Stratified split preserves the original class proportions.
• Balanced sampling changes the stream so classes appear more evenly.

Many teams ask for stratified sampling when what they really want is balanced training. Be clear about the goal before building the pipeline.

Add Shuffle, Batch, and Prefetch Carefully

Once the class logic is correct, the rest of the pipeline looks normal.

train_ds = balanced_ds.shuffle(100).prefetch(tf.data.AUTOTUNE)

Place expensive transforms after you are sure the sampling behavior is correct. That makes debugging much easier, especially when the class mix in the final batches does not match your expectations.

Verify the Output Distribution

Do not assume the dataset is stratified or balanced just because the code looks right. Sample a few hundred labels from the pipeline and count them. That quick check catches missing repeat(), wrong weights, or label filters that silently exclude one class.

1label_counts = {0: 0, 1: 0}
2
3for _, batch_labels in balanced_ds.take(50):
4    for label in batch_labels.numpy():
5        label_counts[int(label)] += 1
6
7print(label_counts)

Alternative: Rejection Resampling

TensorFlow also offers rejection resampling for some workflows, but it is more complex to reason about than explicit class datasets. When the label set is small and well-defined, splitting by class and recombining with explicit weights is usually easier to understand and maintain.

Common Pitfalls

• Mixing up a stratified split with balanced training resampling.
• Filtering by label after heavy preprocessing, which makes the pipeline slower than necessary.
• Forgetting to call repeat() on minority-class datasets when you want a long balanced stream.
• Checking only the first few examples and assuming the sampling ratios are correct without measuring them.

Summary

• Use a stratified split before tf.data when you want training and validation sets with similar class ratios.
• Use sample_from_datasets when you want a balanced training stream from an imbalanced dataset.
• Stratified and balanced are related ideas, but they solve different problems.
• Keep the class-sampling logic simple before adding the rest of the input pipeline.
• Validate the actual class distribution in the output batches instead of assuming the pipeline is correct.