Creating a TF Dataset using a Data Generator

Source : GitHub

In this article, we are going to build a tf.data.Dataset from a data generator. This data generator will perform data fetching, pre-processing, augmentations within itself.

We will be using the array dataset for this article, that is, the data and labels are stored in an array. If the dataset can’t be stored in an array, we can create a DataFrame from the filenames and use it in place of the data array. By changing some methods, this generator can be adapted to any dataset.

One important feature of the generator is that it should be callable. For this, we need the method __call__.

Let us look at a sample code for creating such a Generator

#Instantiate a Generator

dg = ArrayDataset(‘train’, tr_x, tr_y, mean, std, transforms)

mean and std will depend on the dataset and transforms is a function passed to augment the data.

#Instantiate a Dataset

ot = (tf.float32, tf.int64)
os = (tf.TensorShape([tr_x.shape[1:]), tf.TensorShape([]))
ds = tf.data.Dataset.from_generator(dg, 
                                    output_types = ot, 
                                    output_shapes = os)
ds = ds.batch(BATCH_SIZE)

To use this Dataset in training, we can just iterate over it like

for batch in ds:
    x,y = batch
    # Training step

Here is the link to the Colab notebook
https://colab.research.google.com/drive/1pH9NKhAHT40ygOOf4bwld8j221r9SRq_?usp=sharing

Clap if you find this post useful. Comment your feedback, if any.

Originally published on Medium.