For new startups that need to quickly validate their product idea, machine learning has one significant obstacle. The cost of building the dataset and training a new model can be a showstopper. That’s where transfer learning comes in — get the ML functionality needed to prove your concept without breaking the bank.
In this blog post, we will take a real-life machine learning (ML) problem and solve it with the power of transfer learning — achieving 97% image prediction accuracy with 30 seconds of training and just 200 images. Training from scratch would have taken hours, require thousands of images and much more additional work.
Problem & Solution
Problem: Proper Classification
The ResNet50 pre-trained image classification ML model does not support images of ‘beds’. For a product that relies on the correct identification and classification of furniture items, this was a serious shortcoming.
This was one of our problems in the project we outlined here: Image Similarity Search — from Research to Production. Because of this we were unable to support bed images in the MVP phase.
We were using the ResNet50 model for image similarity search and we wanted to use the same model for image classification as well. But ResNet50, which supports 1000+ image classes, did not support furniture class ‘bed’ — simply because it was not trained to do so.
Solution: Transfer Learning
In order to classify ‘bed’ images, the first two options that come to mind are “transfer learning” and “learning from scratch”. We chose transfer learning since it requires less time to train, less data, less engineering time, less code/design complexity and thus less problems — I think I made my point :)
Let’s check the results right away then we will dive into more technical details later on.
The section below is the output of our train & test python script transfer_learn.py. The output is simplified, in the first step we trained our model with transfer learning techniques using only 100 images per class (200 in total, this number is very small since we are using transfer learning).
In the second step, we tested the accuracy of our trained model with 200 images which this model has never seen. Finally, we get 97% test prediction accuracy (and 95% training accuracy) after 30 seconds of training.
(Note: Validation step for training is omitted for simplicity.)
Why Transfer Learning & why is it so important?
With transfer learning the training time can be reduced dramatically — up to 80x. Hours instead of days, minutes/seconds instead of hours, depending on the model, dataset and use case. Other advantages are that the training task, code and model design becomes less complex. These all reduced engineering time and effort.
With transfer learning, the knowledge from a pre-trained model is carried over so we can re-use proven pre-trained models for a similar task.
Let’s see some numbers from Intel — a comparison of transfer learning vs training from scratch (aka traditional training):
What is Transfer Learning?
Transfer learning lets us utilize gained knowledge of an existing model — in our case the ResNet50 image classifier — and apply it to a similar problem.
Transfer learning is basically transferring knowledge (by re-training certain parts of a ML model with a smaller dataset) from one machine learning model to another instead of training a model from scratch. For example, knowledge gained while learning to recognize cars could be transferred to recognize trucks.
First let's show the application of the transfer learning algorithm to categorize images:
Now we can see how transfer learning gives us faster training times with a smaller dataset:
The Code & Our Transfer Learning Strategy
In this section we will summarize our transfer learning strategy and explain its code.
ResNet50: We are using the famous deep neural network ResNet50 as the pre-trained base model for our Transfer Learning. The ResNet50 model has already been trained on ImageNet and is quite a beast — 150 GB, ~1.2 million images for training and validation, 1,000 classes.
But, no beds! So we're going to repurpose it to classify new images of beds and sofas.
Transfer Learning Strategy:
- Load and remove the last predicting layer of the pre-trained ResNet50 model and replace it with our new dense layers.
- Freeze (make non-trainable) the weights of ResNet50 pre-trained model and use it as a feature extractor (Weights of the pre-trained model are frozen and will not be updated during the training).
- Only train the newly added dense layers, which are created from randomly initializing the weights
So, in the end the new model we'll create will consist of a base pre-trained network (from ResNet50) and a new dense network classifier.
The related python code can be seen below and it is heavily commented. Other parts of the code are not shown but can be found in the Github repository linked below.
So, in the end we’ve got the proper classification of beds for our furniture image recognition project. By using the right transfer learning model along with the ResNet50 classifier, we saved a bunch of time without sacrificing accuracy.
Don’t let the estimated cost of model training hold you back from launching an ML-enabled product! As we’ve shown, there are often ways to get the functionality you need without training from scratch.
Thanks for reading 😀
Other Sources and further reading:
- Use Transfer Learning For Efficient Deep Learning Training On Intel®
- Introduction to Transfer Learning with TensorFlow 2.0
- Tutorial Keras: Transfer Learning with ResNet50