Keras Data Generators and How to Use Them
You probably encountered a situation where you try to load a dataset but there is not enough memory in your machine. As the field of machine learning progresses, this problem becomes more and more common. Today this is already one of the challenges in the field of vision where large datasets of images and video files are processed.
Here we will focus on how to build data generators for loading and processing images in Keras.
What is the functionality of the data generator
In Keras Model class, there are three methods that interest us: fit_generator, evaluate_generator, and predict_generator. All three of them require data generator but not all generators are created equally.
Let’s look into what kind of generator each method requires:
fit_generator
Requires two generators, one for the training data and another for validation. Fortunately, both of them should return a tuple (inputs, targets) and both of them can be instance of Sequence class.
evaluate_generator
The data generator here has same requirements as in fit_generator and can be the same as the training generator.
predict_generator
The generator here is a bit different. It should return only inputs.
With that in mind, let’s build some data generators. Because of the similarity between the generator in fit_generator and evaluate_generator, we will focus on building data generators of fit_generator and predict_generator.
ImageDataGenerator class
The ImageDataGenerator class is very useful in image classification. There are several ways to use this generator, depending on the method we use, here we will focus on flow_from_directory takes a path to the directory containing images sorted in sub directories and image augmentation parameters.
Let’s look on an example:
We will use a dataset that can be downloaded from https://www.kaggle.com/c/dogs-vs-cats/data where the structure is as follows:
data/
train/
dogs/
dog001.jpg
dog002.jpg
...
cats/
cat001.jpg
cat002.jpg
...
validation/
dogs/
dog001.jpg
dog002.jpg
...
cats/
cat001.jpg
cat002.jpg
...
First, let’s import all the necessary libraries and create a data generator with some image augmentation.
Finally, create a model and run the fit_generator method.
The ImageDataGenerator is an easy way to load and augment images in batches for image classification tasks. But! What if you have a segmentation task? For that, we need to build a custom data generator.
Flexible data generator
To build a custom data generator, we need to inherit from the Sequence class. Let’s do that and add the parameters we need.
The Sequence class forces us to implement two methods; __len__ and __getitem__. We can also implement the method on_epoch_end if we want the generator to do something after every epoch.
The __len__ method should return the number of batches per epoch. One possible implementation is shown below.
on_epoch_end in this example can shuffle the indexes for the training if shuffle=True. But there can be any logic here that we want to run after every epoch.
The second method that we must implement is __getitem__ and it does exactly what you would expect. It should return a batch of images and masks if we are predicting. This can be controlled by setting to_fit to True or False.
The entire data generator should be similar to this:
Assuming we have two directories, one holds the images and the other holds the mask images and every image has a corresponding mask with the same name, the following code will train the model using the custom data generator.
Finally if we want to make predictions with the data generator, to_fit should be set to False and predict_generator should be called.
Conclusions
While Keras provides data generators, they are limited in their capabilities. One of the reasons is that every task is needs a different data loader. Sometimes every image has one mask and some times several, sometimes the mask is saved as an image and sometimes it encoded, etc…
For every task we will probably need to tweak our data generator but the structure will stay the same.