How to use TensorFlow Hub with code examples
Any deep learning framework in order to be successful, has to provide a good collection of state of the art models, along with its weights trained on popular and widely accepted datasets. Many frameworks usually refer to this collection as Model Zoo. TensorFlow in it’s initial versions provided this model zoo through it’s TF-Slim framework. But TF-Slim’s usability for models was not intuitive and over the course of time, TF-Slim has also been deprecated.
TensorFlow, now has come up with a better framework known as TensorFlow Hub which is very easy to use and is well organised. With TensorFlow Hub, you can confidently perform the widely utilised activity of Transfer Learning by importing large and popular models in a few lines of code. TensorFlow Hub is very flexible and provides the facility to host your models to be used by other users. These models in TensorFlow Hub are referred to as modules. In this article, let us look into basics of how to use module from TensorFlow Hub, it’s various types and code examples.
It is important to note that TensorFlow Hub Module just provides us with graph comprising of architecture of model along with it’s weights trained on certain datasets. Most of the modules give access to internal layers of the model which can be used based on different use cases. However, some of the modules are not fine-tunable. Before starting with development, it is advisable to look into the description provided about the module in TensorFlow Hub website.
Pre-requisite understanding
Next, let us look into some of the important aspects and functions which can be used to understand more details about the TensorFlow Hub’s module.
1) Module Instantiation: Various modules made up with different models (Inception, ResNet, ElMo etc) serving different purposes (image classification, text embeddings etc) are hosted in TensorFlow Hub website. The user has to browse through the catalogue of modules and then once finalised with his purpose and model, needs to copy the URL of the model where it is hosted. Then, the user can instantiate his module like this:
Apart from the URL parameter, the other most notable parameter is ‘trainable’. If user wishes to fine-tune/modify the weights of the model, this parameter has to be set as True.
2) Signature: The signature of the module specifies what is the purpose for which module is being used for. All the module, comes with the ‘default’ signature and makes use of it, if a signature is not explicitly mentioned. For most of the modules, when ‘default’ signature is used, the internal layers of the model is abstracted from the user. The function used to list out all the signature names of the module is get_signature_names().
3) Expected inputs: Each of the module has some set of expected inputs depending upon the signature of the module being used. Though most of the modules have documented the set of expected inputs in TensorFlow Hub website (particularly for ‘default’ signature), some of them haven’t. In this case, it is easier to obtain the expected input along with it’s size and datatype using the get_input_info_dict().
4) Expected outputs: In order to build the remaining part of the graph after the TensorFlow Hub’s model is built, it is necessary to know the expected type of output. get_output_info_dict() function is used for this purpose. Note that for ‘default’ signature, there will be usually only one output, but when you use a non-default signature, multiple layers of the graph will be exposed to you.
5) Collecting required layers of module: After instantiating the module, you have to extract your desired layer/output from module and add it to the graph. These are some of the ways to do it:
6) Initialising TensorFlow Hub operations: The resulting output weight/values of all operations present in modules are hosted by TensorFlow Hub in a tabular format. This needs to be initialized using tf.tables_initializer() along with the initializations of regular variables. The code block looks like this:
Code skeleton block
Once the complete graph comprising of your module, learning algorithm optimisers, objective function, custom layers etc is built, this is how the graph part of your code may look like.
I am making use of two modules. The first module is built using bare minimum code which implicitly uses default signature and layer. In the second module, I am explicitly specifying default signature and layer. In the similar manner, we can specify the non-default signature and layers.
Now that we have looked into how to integrate a module in the graph, let’s look into various types of modules present in TensorFlow Hub at the time of writing this article. (Modules will be continuously added/updated in TensorFlow Hub) In each of the example, I have integrated it with Dataset and Iterator pipeline and saved with Saved_Model. (You can check my detailed Medium post on usage of Dataset and Iterator with code examples)
1) Image Classification Modules
As the name suggests, these set of modules are used for image classification. In this category of modules, the complete architecture of network is given. This includes the final dense layer which is used for classification. In the code example for this category, I am just going to classify the Kaggle’s cat dog classification problem into 1001 ImageNet classes using the Inception V3 module.
This is the set of the output prediction generated for some of the samples.
If you print the signatures of these modules, most of them will be having a total of three signatures which is ‘default’, ‘image_classification’ and ‘image_feature_vector’. Though the signature of ‘image_feature_vector’ appears in this category of modules, this signature will exhibit similar characteristic as that of feature vector category of modules which is explained next.
2) Feature Vector Modules
This is very similar to that of Image Classification Modules with the only difference being that module doesn’t include the final dense classification layer. In the code example for this category, I am going to finetune Resnet-50 module for Hackerearth’s multi-label animal attribute classification problem.
During fine-tuning operation, very often, it becomes necessary to optimise the final layers of the module in addition to your customly appended dense layers. During such scenarios, you will have to find the names of scope or variables you wish to fine-tune. I am attaching the relevant part of code to perform such procedure in Feature Vector Module.
If the nodes in module’s graph are not named intuitively (very unlikely scenario), you will have to figure out the architecture of the module and then map them to the nodes being printed.
In the code example, I have fine-tuned the final few layers’ variables falling under the variable scope of ‘resnet_v2_50/block4’. By training for 10 epochs, you will get a F1 score of 0.95819 effortlessly for this problem.
3) Video Classification Modules
Video classification refers to categorising the nature of activities happening in the video clip. In the code example, I have trained the 20BN-Jester dataset comprising of 27 classes on Inception 3D module. Again, with just 3 epochs of training, I have achieved a score of 91.45% accuracy. The relevant part of module’s code is shown below.
In order to test the accuracy of the model, I built up a simple real-time video classification application and here is a small video of it.
4) Text Embedding Modules
TensorFlow Hub has supplied several text embedding modules in not only English but also in several other languages like German, Korean, Japanese, etc. (At the time of writing this article), many of them produce embeddings at the sentence level and not at the word level. In the code example, I have worked on Kaggle’s sentiment analysis on movie review. Using the ELMO embeddings module, I trained the network using embeddings produced at sentence level and word level separately. The network comprised of single layer of bi-directional LSTM. At the sentence level, after training for 6 epochs, the accuracy achieved was 65%. Attaching the relevant part of module’s code.
5) Image Augmentation Modules
Image augmentation is an essential component in the training pipeline to increase the accuracy of the model. None of the image augmentation modules (at the time of writing) have any variables in them, and consequently, these modules are not fine-tunable/trainable. Using the signature of ‘from_decoded_images’, you can directly feed the images to these module. An example code block along with the generated augmented images is shown below.
The real utility of image augmentation comes when you club this module with other modules from image classification, feature vector, etc in your pipeline. In the previously trained multi-label animal attribute classification problem using Feature Vector module, I have added crop_color image augmentation module, and with the same number of epochs, my F1 score improves to 0.96244. In this code example, I make use of default signature, which will expect the image in encoded form as its input. Now, as my image augmentation module resizes the input image to a fixed size for further processing by Resnet50 module, I have saved the work of explicitly resizing the images myself which I had done previously.
You can also check my previously written article on image augmentation covering the basic types of augmentations.
6) Object Detection Modules
Object detection modules do not support fine-tuning, so you will have to perform training from scratch if you have your own dataset. Batching of the data is also not supported at the moment. In the code example, I will be only performing inference on the images using FasterRCNN on Inception-ResNet V2 module. I have enclosed the module part of code and images generated from the module below.
7) Generator Modules
These correspond to Generative Adversial Networks(GAN). Some of the modules have not exposed their Discriminator part of the network. Just like object detection, even here in code sample, I will be just performing inference. Using the Progressive GANs module which is trained on CelebA dataset, I will be generating new faces. I have enclosed the module part of code and images generated from the module below.
If you would like to see the complete code integrated with Dataset and Iterators, and saved with Saved_Model, you can check out the code in my GitHub repository.
If you would like to read more about Dataset and Iterators of TensorFlow along with code examples, you can read about them in my another Medium post.
Lastly, do let me know your thoughts about this article through comments section.