How to Divide Datasets Into Mini-Batches in Python

In machine learning and deep learning, processing an entire dataset at once is often impossible due to memory constraints. Instead, we use mini-batches, i.e. small, manageable chunks of data processed sequentially.

This guide explains how to implement a memory-efficient data pipeline in Python using generators to slice a dataset into mini-batches.

Understanding Generators and Slicing

To divide a list into chunks, we rely on two Python features:

1. List Slicing: list[start:end] extracts a specific portion of the data. Python handles out-of-bounds indices gracefully (e.g., if you ask for 5 items but only 3 remain, it simply returns the 3).
2. Generators (yield): Instead of creating a list of lists containing all batches (which doubles memory usage), we use yield to return one batch at a time.

Step 1: Implementing the Data Pipeline

We will create a function data_pipeline that accepts raw data and a batch size.

Memory Inefficient vs. Efficient Approach

def bad_pipeline(data, batch_size):
    batches = []
    for i in range(0, len(data), batch_size):
        batches.append(data[i : i + batch_size])
    return batches

from typing import Generator, List

def data_pipeline(data: List[List[int]], batch_size: int) -> Generator[List[List[int]], None, None]:
    """
    Yields mini-batches from the dataset sequentially.
    """
    for i in range(0, len(data), batch_size):
        batch = data[i : i + batch_size]
        yield batch

Step 2: Testing the Mini-Batch Generator

To verify the logic, we will iterate through the generator. We expect the final batch to contain fewer items if the total count isn't perfectly divisible by the batch size.

if __name__ == "__main__":
    dataset = [
        [1, 2], 
        [1, 3], 
        [3, 5], 
        [2, 1], 
        [3, 3]
    ]
    
    BATCH_SIZE = 2
    
    print(f"Total items: {len(dataset)}")
    print(f"Batch size:  {BATCH_SIZE}")
    print("-" * 20)

    batch_generator = data_pipeline(dataset, BATCH_SIZE)

    for i, batch in enumerate(batch_generator):
        print(f"Batch {i+1}: {batch}")

Output Validation:

Total items: 5
Batch size:  2
--------------------
Batch 1: [[1, 2], [1, 3]]
Batch 2: [[3, 5], [2, 1]]
Batch 3: [[3, 3]]

Conclusion

Creating a mini-batch generator is a fundamental skill for data processing.

1. Use range(0, len(data), batch_size) to calculate start indices.
2. Use data[i : i + batch_size] to slice the data safely.
3. Use yield to create a memory-efficient iterator.