Data Analysis Pipeline with TensorFlow Data: tf.data | Towards Data Science

While training a neural network, it is quite common to use ImageDataGenerator class to generate batches of tensor image data with real-time data augmentation. However, in this post, I will discuss tf.data API, using which we can build a faster input data pipeline with reusable pieces. As mentioned in the TensorFlow documentation —
The tf.data API makes it possible to handle large amounts of data, read from different data formats, and perform complex transformations.
This post will be short and crisp including a working example using tf.data; so without any delay let’s begin. Codes used here can be found in my notebook and the link is given at the end of the post.
The main component of tf.data that we will use here to build an efficient pipeline is tf.data.Dataset API. It can be used for the following jobs —
Create ‘dataset’ object from input data.Apply data-set transformations for preprocessing.Iterate over the dataset in a streaming fashion and process the elements.
Create ‘dataset’ object from input data.
Apply data-set transformations for preprocessing.
Iterate over the dataset in a streaming fashion and process the elements.
We will go through each of these steps using an example to understand and implement tf.data . First, we will create ‘dataset’ objects using from_tensor_slices method.
To Understand how from_tensor_slices method works, let’s get started by loading the CIFAR-10 data1.
import matplotlib.pyplot as pltimport numpy as npimport timeimport tensorflow as tffrom tensorflow.keras.datasets import cifar10(x_train, y_train), (x_test, y_test) = cifar10.load_data()print ('check shapes: ', x_train.shape, y_train.shape, x_test.shape, y_test.shape)>>> check shapes:  (50000, 32, 32, 3) (50000, 1) (10000, 32, 32, 3) (10000, 1)
We turn the labels to categorical representation by using to_categorical
train_lab_categorical = tf.keras.utils.to_categorical(y_train, num_classes=10, dtype=’uint8')test_lab_categorical = tf.keras.utils.to_categorical(y_test, num_classes=10, dtype=’uint8')from sklearn.model_selection import train_test_splittrain_im, valid_im, train_lab, valid_lab = train_test_split(x_train, train_lab_categorical, test_size=0.20, stratify=train_lab_categorical,random_state=40, shuffle = True)print ("validation labels shape: ", valid_lab.shape)>>> validation labels shape:  (10000, 10)
We will now combine images and labels to create ‘Dataset’ objects as below —
training_data = tf.data.Dataset.from_tensor_slices((train_im, train_lab))validation_data = tf.data.Dataset.from_tensor_slices((valid_im, valid_lab))test_data = tf.data.Dataset.from_tensor_slices((x_test, test_lab_categorical))print (‘check types; '\n' ‘, type(training_data)) >>> check types;  <class 'tensorflow.python.data.ops.dataset_ops.TensorSliceDataset'>
From type we cannot get much but let’s try to figure out what are these ‘Dataset’ objects and how we can use them. First, we check the type specification of an element of this dataset using element.spec() .
print (training_data.element_spec)>>> (TensorSpec(shape=(32, 32, 3), dtype=tf.float64, name=None), TensorSpec(shape=(10,), dtype=tf.uint8, name=None))
This is more clear to understand. We see that from_tensor_slices preserves the structure of the input tensors. An element of this dataset consists of Image (with shape: 32, 32, 3) and corresponding label (with shape: 10,). Our next question should be how to access the elements within the ‘Dataset’ object? We can create an iterator object and access the elements by using next as below —
train_iter_im, train_iter_label = next(iter(training_data))print (train_iter_im.numpy().shape, train_iter_label.numpy().shape)>>> (32, 32, 3) (10,)
Instead of next, iterand .numpy() , we can return an iterator that converts all elements of the dataset to numpy array in a single step as below —
train_iter_im1, train_iter_label1 = next(training_data.as_numpy_iterator())print (train_iter_im1.shape, train_iter_label1.shape)>>> (32, 32, 3) (10,)
Let’s visualize some training images and corresponding labels as below —
class_types = [‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’,‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’] # from cifar-10 websitecheck_list = list(training_data.as_numpy_iterator())fig = plt.figure(figsize=(10,10))for i in range(12):  plt.subplot(4,3,i+1)  plt.xticks([])  plt.yticks([])  plt.grid(False)  plt.imshow(check_list[i][0], cmap='gray')  plt.xlabel(class_types [np.argmax(check_list[i][1])], fontsize=13)plt.tight_layout()plt.show()
Generating training and validation batches with tf.data is way faster than ImageDataGenerator. Let’s compare; First, we use ImageDataGenerator without using any augmentation —
from tensorflow.keras.preprocessing.image import ImageDataGeneratorst1 = time.time()train_datagen = ImageDataGenerator().flow(x=train_im, y=train_lab, batch_size=128, shuffle=True,)et1 = time.time()print (‘time taken: ‘, et1-st1)>>> time taken:  0.2080059051513672
Now we will use tf.data —
autotune = tf.data.AUTOTUNEst2 = time.time()train_datagen_tfdata = training_data.shuffle(buffer_size=40000).batch(128).prefetch(buffer_size=autotune)et2 = time.time()print (‘time taken: ‘, et2-st2) >>> time taken:  0.006140947341918945
So tf.data is about ~34 times faster than ImageDataGenerator and one of the main reasons for that is a technique called ‘Prefetching’. We apply this prefetching while generating batches using tf.data via prefetch(buffer_size=tf.data.AUTOTUNE) . In the TensorFlow documentation series of examples are given with excellent explanations but here in brief what is happening. The data pipeline can be thought of as a combination of ‘producer’ (generating batches) and ‘consumer’ (batches that are used to train the neural net) and prefetch transformation provides benefits whenever there is an opportunity to overlap the work of a ‘producer’ with the work of a ‘consumer’.
Since we are not using ImageDataGenerator, we will use Keras pre-processing layers to augment images. Apart from the generic augmentation, we can also rescale the images (pixel values between 0 and 1) which is a necessary step for the problem at hand. The idea is — we will use keras.sequential just as if we are building a deep conv-net, and instead of Convolution or Pooling layers, we will use augmentation layers. Finally, we will visualize one of the examples of augmented images and below is the code block. More examples are given in the TensorFlow documentation.
We will use the augmentation layers within Sequential as a single layer in our model. This also helps to train the model faster since the data augmentation is now going on in the GPU and not in CPU. Also, to keep things simple, instead of building a model from a scratch we will use pre-trained ResNet with ImageNet data. I described building a ResNet from scratch in a separate post. For a multi-class problem like CIFAR-10, our loss function will be categorical cross-entropy. The code block below combines all of these together. Let’s see —
With 60 epochs, I obtained the training and validation plots as below —
We can also plot the confusion matrix for the test data and below is the plot —
To conclude, we went through an example of using tf.data to create a faster, lucid and reusable pipeline and for model deployment in the real world, this is generally a better option than ImageDataGenerator. ImageDataGenerator is a great option to get started with but, tf.data can autotune the process of generating batches and training simultaneously, depending on the hardware/experiment design and the complete process is way faster.
Stay strong and cheers!!
Codes used here are from My Notebook.
CIFAR-10 Data Source and Citation:
[1] ‘Learning multiple layers of features from tiny images’: Alex Krizhevsky; 2009. Source of CIFAR-10.