Skip to content
Tutorial Reference

How to Implement a Balanced Batch Generator in Python

Learn how to create a custom Python batch generator that dynamically balances class distribution in machine learning datasets using upsampling and downsampling techniques.

In machine learning, training a model on an imbalanced dataset, where one class has significantly more samples than others, can lead to biased predictions. A common strategy to mitigate this is to create a data pipeline that dynamically generates batches with a balanced distribution of classes.

This guide explains how to implement a custom batch generator in Python that performs upsampling and downsampling on the fly to ensure equal class representation in every training step.

Understanding the Strategy

To balance a batch of size N across K unique classes:

  1. Group the dataset samples by their labels.
  2. Calculate quota: Determine how many samples (Q) each class should contribute to the batch (Q = N // K).
  3. Sample: Randomly select Q samples from each class.
    • If a class has more samples than Q, we downsample (pick a subset).
    • If a class has fewer samples (and we allow replacement), we upsample (repeat samples).
  4. Fill Remainder: If N isn't perfectly divisible by K, randomly fill the remaining slots.

Step 1: Grouping Data by Label

First, we need to organize our raw data so that we can access all samples belonging to a specific class instantly. We use Python's collections.defaultdict.

from collections import defaultdict


def group_data_by_label(data):
    counter = defaultdict(list)
    for x, y in data:
        counter[tuple(y)].append(x)
    return counter

Step 2: Calculating Batch Distribution

Given a batch size, we calculate how many samples each class must provide (pre_num) and how many slots remain (num_left) if the division isn't exact.

batch_size = 6
num_classes = 2

pre_num = batch_size // num_classes  # 3

num_left = batch_size % num_classes  # 0

Step 3: Implementing the Sampling Logic

We iterate through our grouped data and select samples.

import random

def sample_balanced_batch(counter, pre_num, num_left):
    batch_data = []
    
    for label_tuple, features_list in counter.items():
        samples = random.sample(features_list, pre_num)
        
        batch_data.extend([[sample, list(label_tuple)] for sample in samples])
        
    all_labels = list(counter.keys())
    for _ in range(num_left):
        y = random.choice(all_labels)
        x = random.choice(counter[y])
        batch_data.append([x, list(y)])
        
    return batch_data

Complete Code Implementation

Here is the complete, runnable pipeline combining all steps.

import random
from collections import defaultdict

def unbalanced_data_pipeline(data, batch_size):
    """
    Generates a batch of data with balanced class distribution.
    """
    counter = defaultdict(list)
    for x, y in data:
        counter[tuple(y)].append(x)
    
    num_classes = len(counter.keys())
    pre_num = batch_size // num_classes
    num_left = batch_size % num_classes
    
    batch_data = []
    
    for y, x in counter.items():
        if len(x) >= pre_num:
            samples = random.sample(x, pre_num)
        else:
            samples = random.choices(x, k=pre_num)
            
        batch_data.extend([[sample, list(y)] for sample in samples])
    
    keys_list = list(counter.keys())
    for _ in range(num_left):
        y = random.choice(keys_list)
        x = random.choice(counter[y])
        batch_data.append([x, list(y)])
        
    return batch_data

if __name__ == "__main__":
    data = [
        [[1, 2, 5], [1, 0]], [[1, 6, 0], [1, 0]], [[4, 1, 8], [1, 0]], # 3 samples of Class A
        [[7, 0, 4], [0, 1]], [[5, 9, 4], [0, 1]], [[2, 0, 1], [0, 1]], # 9 samples of Class B
        [[1, 9, 3], [0, 1]], [[5, 5, 5], [0, 1]], [[8, 4, 0], [0, 1]],
        [[9, 6, 3], [0, 1]], [[7, 7, 0], [0, 1]], [[0, 3, 4], [0, 1]],
    ]
    
    print("Generating balanced batch of size 6...")
    batch = unbalanced_data_pipeline(data, 6)
    
    class_a_count = sum(1 for item in batch if item[1] == [1, 0])
    class_b_count = sum(1 for item in batch if item[1] == [0, 1])
    
    print(f"Batch content: {batch}")
    print(f"Class [1, 0] count: {class_a_count}")
    print(f"Class [0, 1] count: {class_b_count}")

Output:

Generating balanced batch of size 6...
Batch content: [[[1, 2, 5], [1, 0]], [[4, 1, 8], [1, 0]], [[1, 6, 0], [1, 0]], [[5, 5, 5], [0, 1]], [[1, 9, 3], [0, 1]], [[8, 4, 0], [0, 1]]]
Class [1, 0] count: 3
Class [0, 1] count: 3

Conclusion

By manually controlling the batch generation process, you can ensure that your machine learning model sees a balanced representation of all classes during training.

  1. Map your data to separate buckets based on labels.
  2. Calculate the exact number of samples needed per class for your batch size.
  3. Sample randomly from these buckets to fill the batch.

This technique is a straightforward and effective way to handle imbalanced datasets without needing external libraries.

Navigation