How can I become a TPU expert?

Another post starts with you beautiful people!

I have two good news for all of you! First good news is that Tensorflow has released it's new version (TF 2.1) which is focused on TPUs and the most interesting thing about this release is that it now also supports Keras high level API. And second wonderful news is to help us get started Kaggle has launched a TPU Playground Challenge. This means there is no any way to stop you learning & using TPUs. In this post I am going to share you how to configure and use TPUs while solving a image classification problem.

What are TPUs?

You must have heard about TPU while using Google Colab. Now Kaggle also supports this hardware accelerator. TPUs or Tensor Processing Units are hardware accelerators specialized in deep learning tasks. They were created by Google and have been behind many cutting edge results in machine learning research. Kaggle Notebooks are configured with TPU v3-8s, which is a specialized hardware with 4 dual-core TPU chips for a total of 8 TPU cores. We can enable this accelerator in our Kaggle notebook by selecting  'TPU v3-8' option in Accelerator menu. We can use up to 30 hours per week of TPUs and up to 3 hours at a time in a single session.

How to enable TPU training in Tensorflow Keras?

Once we enable the TPU accelerator, next step is to enable the TPU training with Keras. You just need to first import the Tensorflow library and execute the following code to enable this-

Here TPUClusterResolver() locate the TPUs, TPUStrategy object contains the necessary distributed training code that will work on TPUs with their 8 compute cores. Finally, we use the TPUStrategy by instantiating our model in the scope of the strategy. This creates the model on the TPU. Model size is constrained by the TPU RAM only, not by the amount of memory available on the VM running our Python code. Model creation and model training use the usual Keras APIs.

TPUs and GCS?

TPUs read training data exclusively from Google Cloud Storage (GCS). Format of this training data in GCS is TFRecords. The TFRecord format is a container format frequently used in Tensorflow to group and shared data files for optimal training performance. Each file contains the id, label (the class of the sample, for training data) and img (the actual pixels in array form) information for many images. On Kaggle we can also use KaggleDatasets().get_gcs_path() to obtain the GCS path of a Kaggle dataset. 

How to load a dataset from TFRecords file?

We can load a dataset from TFRecords files by using following code snippet:-

From multiple TFRecords we can also enable parallel streaming with below code snippet:

Above three code snippets are only required to use TPUs. Now I will show you how can we use TPUs while solving a 100 flower classification problem. I have tried different architectures with pretrained weights- VGG16, VGG19, ResNet, DenseNet etc and found best result (0.92313) with DenseNet201 with default learning rate. With a custom learning rate function I even reached a public score of 0.93205. In the end of the post I will share my Kaggle notebook link so that you can check it and also fork it to explore with any new model or settings. Let's start the steps I have followed to solve the problem-

1. Initiate and enable TPUs:

2. Load the Datasets: In this playground challenge, we're classifying 104 types of flowers based on their images drawn from five different public datasets and images are provided in TFRecord format. The dataset is already divided into train, validation and test sets. Each set contains 4 types of images having dimension of 192x192, 224x224, 331x331 and 512x512. We will load the datasets from GCS in following way:

Here I am using 224x224 dimension images, You can change it to any of the given dimension. After selecting the dataset we can use Tensorflow's io.gfile.glob() function to load the datasets. In above way you can load any TFRecord format files. Next step is to wrap your model initialization and compilation code in the TPU strategy scope. In my initial try I experimented with various pretrained weights of Keras API. I am here showing the best one.

3. Compile model with TPU: Model initialization, compiling and training step is same as you did while using keras API. Only different point here is to wrap our code inside TPU strategy scope. Below is code snippet of the same I followed:

4. Train the model: Training of the model is as same as we did with any of Keras API. We need to call .fit() function with required parameters. below is code snippet of the same:

Once my model training was completed. I calculated the confusion matrix and got below result:

When I submit my final submission file I got 0.93205 public score while the top score is 0.96652 which means there are many improvement still left. I am still trying. What about you? You don't need to type the above mentioned code manually. Just click on 'Copy and Edit' button in my notebook: my notebook url and start experimenting with new things till you get the maximum score. Till then Go chase your dreams, have an awesome day, make every second count and see you later in my next post.