How to Choose the Best Keras Pre-Trained Model for Image Classification | by Marie Stephen Leo | Towards Data Science

Deciding which pre-trained model to use in your Deep Learning task ranks at the same level of classical dilemmas like what movie to watch on Netflix and what cereal to buy at the supermarket (P.S. buy the one with the least sugar and highest fiber content). This post will use a data-driven approach in Python to find out the best Keras Pre-Trained model for the cats_vs_dogs dataset. This post and the code provided will also help you easily choose the best Pre-Trained model for your problem’s dataset.
BackgroundCriteria for selecting modelsCodeResources
Background
Criteria for selecting models
Code
Resources
Transfer learning is a technique in machine learning where we can take a model developed on one task and use it as a starting point in some other similar but different tasks. Transfer learning is extremely popular in Deep Learning since the “transfer” of knowledge from one “parent” model to a “child” model means that the “child” model can be trained to high accuracies with a much smaller dataset compared to the “parent” model. You can find more details about Transfer Learning at this excellent [Medium post]
High-level Deep Learning frameworks like TensorFlow and Pytorch have made it incredibly easy to leverage the power of Transfer learning by including several pre-trained models within the package itself. For example, the [TensorFlow Keras API] includes 18 highly advanced model architectures pre-trained on the [“ImageNet”] dataset. You can think of “ImageNet” like the SAT of image classification datasets.
Most Image Classification Deep Learning tasks today will start by downloading one of these 18 pre-trained models, modify the model slightly to suit the task on hand, and train only the custom modifications while freezing the layers in the pre-trained model. This approach gives very high accuracy on real-world image classification tasks since “ImageNet” collects many real-world images.
However, choosing which of the 18 pre-trained models to use is not always an exact science. Many developers stick to the models they are familiar with and gave them good results in the past. As a Data Scientist, I wondered whether there is a better, data-driven approach to selecting a pre-trained model, free from personal biases. Indeed there is! The rest of this post talks about my process with code you can directly use in your projects.
In general, there are two competing criteria while doing any machine learning task in the industry:
Accuracy of the model: Higher BetterSpeed of Model Training and Predictions: Faster Better
Accuracy of the model: Higher Better
Speed of Model Training and Predictions: Faster Better
The two criteria are pretty straightforward. We want the model that gives us the highest accuracy on validation data since it can make useful predictions. We also want the model to train and predict as fast as possible because, in production, we might need to serve hundreds or thousands of predictions every second.
Simple enough? Well, like everything in life, we can’t have our cake and eat it too (hence the title image)! In general, to get higher accuracy, we need to use a “deeper” or larger model. But a larger model has many more parameters that make it slower to execute. We can see the tradeoff between the accuracy and size/number of operations of the model from the below graph [Reference].
So, this is a balance we need to strike. For most of the Deep Learning deployments in the industry, we choose the smallest model that provides us with good enough accuracy. Thus, we need to experiment among many available models to pick the one that meets these criteria.
On a side note, if you only care about achieving the highest accuracy no matter the impact to speed, then your approach could be to combine all these models using ensembling techniques! Indeed, ensembling is quite popular in academia and ML competitions.
My approach is straightforward, as shown in the image below. We will use Python to find all the pre-trained models in Keras and then loop over them one by one. We’ll train the models on the TensorFlow [cats_vs_dogs] dataset in this post. You can replace it with any other dataset, including your own custom dataset.
The first step is to import the necessary packages. If you lack these packages, you can install them using pip install <package-name> . We’ll also set a batch_size for subsequent model training. Since I’m using a low-cost GPU, I’ll use a small batch_size of 32 only.
Automatically get a list of all available pre-trained models from Keras by listing all the functions inside tf.keras.applications. Since each model is instantiated by calling a function from tf.keras.applications, when we list all the functions within this module using inspect.isfunction, we get a list of all the models. As of this writing, in [Tensorflow 2.2], there are a total of 18 models that we can use.
You can add any other custom pre-trained models into the experiments by manually adding elements into the model_dictionary as well by following the syntax model_dictionary["new_model_name"] = new_model_function(), where new_model_function() should return the custom pre-trained model without the final output Dense layer. The model’s input shape should be (224,224,3) if you don’t want to change any other code below.
Download some images to run the experiment. You should modify this step to load your data when you run the experiment for your use case.
Output:
Num train images: 16283         Num validation images: 6979         Num classes: 2         Num iterations per epoch: 508
Preprocess the input images. Some pre-trained models require images to be of size (224,224,3) while some require (331,331,3). We use this step to prepare both sizes of images and normalize them by dividing each pixel by 255. We also one-hot encode the labels so that we can use categorical_crossentropy loss during training.
Loop over each model by downloading the pre-trained model without the output layers and freeze the weights. We then construct an empty Sequential model and first add the pre-trained model to it. We then add a single output Dense layer with softmax activation and compile it with categorical_crossentropy loss. Finally, we train the model by calling model.fit for 3 epochs. We log the number of parameters (size of the model) and each model’s final accuracy into a dictionary to visualize the results. If you get the following error during the model download AttributeError: 'str' object has no attribute 'decode', install a lower version of h5py [source]
pip uninstall h5pypip install h5py<'3.0.0'
The whole process takes ~3 hours on an n1-standard-4(4 vCPUs, 15 GB RAM) with one NVIDIA Tesla T4 GPU on Google Cloud Platform.
Let’s visualize the results by converting it to a DataFrame and sorting in ascending order of num_model_params since our goal is to select the smallest model with good enough accuracy. In this example, the MobileNet model already provides 97% accuracy; thus, we can directly use it. MobileNetV2 and NASNetMobile are two other models we can consider for fine-tuning experiments. If our selected model's accuracy is still insufficient for our task, we can experiment by fine-tuning these selected models further, adding data augmentation, etc. We will apply the typical Deep Learning model improvement experiments on a model that is already proven to be close to our requirements.
Output:
| model_name        | num_model_params | validation_accuracy ||-------------------|------------------|---------------------|| MobileNetV2       | 2257984          | 0.9475569725036621  || MobileNet         | 3228864          | 0.9773606657981873  || NASNetMobile      | 4269716          | 0.9753546118736267  || DenseNet121       | 7037504          | 0.9273535013198853  || DenseNet169       | 12642880         | 0.95572429895401    || VGG16             | 14714688         | 0.9107322096824646  || DenseNet201       | 18321984         | 0.9419687390327454  || VGG19             | 20024384         | 0.8948273658752441  || Xception          | 20861480         | 0.9550078511238098  || InceptionV3       | 21802784         | 0.9859578609466553  || ResNet50V2        | 23564800         | 0.9802263975143433  || ResNet50          | 23587712         | 0.49620288610458374 || ResNet101V2       | 42626560         | 0.9878206253051758  || ResNet101         | 42658176         | 0.49620288610458374 || InceptionResNetV2 | 54336736         | 0.9885370135307312  || ResNet152V2       | 58331648         | 0.9840951561927795  || ResNet152         | 58370944         | 0.49620288610458374 || NASNetLarge       | 84916818         | 0.9795099496841431  |
Plot the results
In conclusion, I’ve presented a data-driven approach to choose the most suitable model from among a group of pre-trained models in the Tensorflow Keras API. In most industrial applications, we would select the model with the least number of parameters that provided good enough accuracy for further exploration. We can easily extend this approach to include other models not offered by Keras by manually adding items into the model_dictionary in [Step 2].
Finally, extensive experimentation and data-driven decision making are the keys to success in all Machine Learning applications. I hope this post stimulates you to think of ways to make more data-driven decisions in your daily work.
You can find all the code used in this post at [GitHub].
You can also use Google Colab to run the code without any setup directly. Just remember to go to Edit → Notebook Settings and select a GPU Hardware Accelerator: [Open In Colab]
Thank you for reading!
Edit: Updated code snippets to GitHub gists