Produce balanced mini batch with Dataset API
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
In machine learning and deep learning tasks, handling imbalanced datasets is a prevalent challenge. An imbalanced dataset can bias a model towards the majority class, leading to suboptimal performance for minority classes. One effective strategy to address this issue is using balanced mini-batches during training. TensorFlow's `tf.data` API provides powerful tools to efficiently manipulate and prepare datasets. This article delves into the implementation of balanced mini-batches using the TensorFlow Dataset API.
Understanding Imbalanced Datasets
An imbalanced dataset occurs when classes are not represented equally. For example, in a binary classification problem, an imbalanced dataset might contain 90% of samples from class 0 and only 10% from class 1. Training a model on such data without adjustments tends to result in high accuracy for the majority class but poor performance on the minority class.
Importance of Balanced Mini-Batches
Balanced mini-batches ensure that each batch of data sent to the model during training contains an equal or nearly equal representation of each class. This can help the model learn equally from all classes and improve its generalization performance on unseen data.
Creating a Balanced Mini-Batch Using TensorFlow Dataset API
To create balanced mini-batches, we can use `tf.data.Dataset` for efficient data input pipelines. With custom sampling strategies, we can ensure that each batch has a balanced representation of each class. Below, we'll walk through an example using TensorFlow.
Step-by-Step Implementation
- Load and Prepare the Dataset:
- Start by loading your data into a `tf.data.Dataset` object. If your data is in NumPy arrays, you can convert it using `tf.data.Dataset.from_tensor_slices`.
- Partition the dataset into two datasets based on class labels.
- Ensure each class dataset's size is a small multiple of your batch size. This can be achieved by repeating the minority class dataset and shuffling it to balance it against the majority class dataset.
- Interleave batches from both datasets to form a balanced dataset and shuffle for randomness.
- Finally, batch the dataset.
- Batch Size: Adjust the batch size according to your dataset size, ensuring it's a multiple that can be handled by both classes.
- Shuffle Buffer: Set a shuffle buffer large enough to randomize your batches meaningfully without missing the balance.
- Performance: Be mindful of computational and memory performance by profiling how the dataset pipeline performs, particularly for large datasets.
Related reading
- Profiling python-tensorflow-1.14
- Proper way to iterate tf.data.Dataset in session for 2.0
- Properly set up exponential decay of learning rate in tensorflow
- Pros and Cons of Amazon SageMaker VS. Amazon EMR, for deploying TensorFlow-based deep learning models?
- Producing a confusion matrix with cross_validate
- Production architecture for big data real time machine learning application?
- Programmatically determine the relative popularities of a list of items books, songs, movies, etc
- Progress indicator during pandas operations

System Design Fundamentals
Build a strong foundation in designing scalable, reliable distributed systems.
View the courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.