Unleashing the Power of Transfer Learning with TensorFlow
Deep learning has revolutionized the field of machine learning in recent years, enabling computers to achieve superhuman performance on tasks like image classification, object detection, natural language processing, and more. However, the Achilles heel of deep learning is often the vast amounts of high-quality labeled training data required to train accurate models from scratch. For many domains, acquiring such extensive datasets is extremely difficult, expensive, or even impossible.
Fortunately, transfer learning provides an innovative solution to this challenge. With transfer learning, instead of training a deep learning model from scratch, you start with a model that has already been pre-trained on a large dataset, usually for a related task. You can then repurpose this model for your own task by retraining just the final layer or fine-tuning some of the deeper layers on your smaller dataset. This allows the model to leverage the robust hierarchical features learned from the initial training, while adapting them to be specific to your target task. The result is often significantly higher accuracy than training the same model architecture from scratch on the limited data.
Transfer learning first rose to prominence in the early 2010s with the remarkable success of deep convolutional neural networks (CNNs) pre-trained on the large-scale ImageNet dataset in computer vision tasks. In 2012, the AlexNet model achieved a top-5 error rate of 15.3% on the ImageNet Large Scale Visual Recognition Challenge (ILSVRC), more than 10 percentage points lower than the previous state-of-the-art [1]. This breakthrough demonstrated the power of deep CNNs trained on massive datasets to learn general visual features applicable across a wide range of recognition tasks. In the ensuing years, increasingly sophisticated CNN architectures like VGGNet, Inception, ResNet, and EfficientNet pushed ImageNet accuracy to new heights while also serving as the foundation for transfer learning pipelines in everything from facial recognition to tumor classification [2].
Beyond computer vision, transfer learning has emerged as a powerful tool in natural language processing thanks to large pre-trained language models like BERT, GPT, XLNet, and T5 [3]. By pre-training on massive unsupervised text corpora to build rich linguistic representations, then fine-tuning on more limited supervised datasets for downstream tasks, these models have achieved record-breaking results in machine translation, question answering, text summarization, and more. Transfer learning is also making inroads in other data-scarce domains like speech recognition, robotics, and bioinformatics.
The efficiency gains of transfer learning can be staggering, as shown in the figure below from a 2014 study comparing CNNs trained from scratch vs. pre-trained CNNs fine-tuned for various computer vision tasks [4]:
For some tasks, the pre-trained models nearly matched the accuracy of the models trained from scratch with 60 times less data! This is the power of transfer learning to combat the data bottleneck.
Transfer learning has proven to be a game changer for many data-scarce applications, perhaps most notably in medicine. Despite widespread digitization of medical data, labeled datasets of medical images, electronic health records, genomic data, etc. remain extremely limited due to privacy concerns, high annotation costs, and the sheer diversity of conditions, modalities and formats. Nevertheless, transfer learning with deep CNNs pre-trained on huge natural image datasets like ImageNet has achieved state-of-the-art results on medical imaging tasks from skin cancer classification to brain tumor segmentation with shockingly small training sets (sometimes just a few hundred images) [5].
Similar success stories have played out in other domains like wildlife conservation, where pre-trained CNNs fine-tuned on limited aerial imagery can accurately estimate animal population sizes [6], and industrial inspection, where models pre-trained on ImageNet can detect manufacturing defects with far less training data than building dedicated models from scratch [7].
In this post, we‘ll dive into the nuts and bolts of how YOU can harness the power of transfer learning using the popular TensorFlow library. TensorFlow is an end-to-end open-source platform for machine learning developed by Google that allows you to easily build, train, and deploy ML models. It has particular strength in deep learning, with a wide variety of pre-trained models and tools to enable transfer learning. We‘ll be using TensorFlow with the Keras API, which provides an intuitive high-level interface for defining and training custom deep learning models. While the standalone Keras package is another option, tf.keras is tightly integrated with TensorFlow, enabling access to the full range of TensorFlow functionalities and ensuring optimal performance on TensorFlow runtime environments.
But before we jump into code, let‘s examine how transfer learning works under the hood. Consider a deep CNN trained for ImageNet classification like the ResNet architecture depicted below:
As the network is trained on the 1.2 million labeled ImageNet examples, it learns a hierarchy of features in its layers [8]. The initial layers learn simple visual building blocks like edges and textures, while deeper layers compose these into more complex concepts like shapes, objects and scenes. Importantly, many of these features, especially in the earlier layers, are not specific to the 1000 ImageNet classes but rather fundamental properties of natural images. ResNet models in particular use skip connections to make feature reuse easier and enable training of very deep networks (hundreds of layers).
These general features are the key to effective transfer learning. By borrowing the weights of these pretrained layers and adapting only a lightweight set of task-specific layers, we can achieve high accuracy on a new task with minimal labeled data and training time. There are several common approaches [9]:
-
Feature extraction – freeze all layers of the pre-trained model and only train a new final layer for the new task. This is the most efficient approach when you have very little data and your new task is similar to the original training task.
-
Fine-tuning – unfreeze some of the deeper pre-trained layers and retrain them along with the new final layer. This allows more customization to the new task but requires more data and risks overfitting. Often combined with gradual unfreezing, where you start training only the final layer, then progressively unfreeze and retrain the deeper layers.
-
Pre-trained weights as initialization – initialize your new model with the pre-trained weights instead of random initialization but retrain the entire model on the new task. Appropriate when you have a lot of data and your new task differs significantly from the original.
-
Multi-task learning – share a base set of layers across related tasks and jointly train task-specific final layers for each one. Allows knowledge transfer across the tasks to improve generalization.
To make things concrete, let‘s consider a case study of binary image classification – the task of predicting which of two classes an input image belongs to. We‘ll use a fun dataset from Kaggle called "Hot Dog – Not Hot Dog" which contains about 500 training images and 500 test images of, you guessed it, hot dogs and things that are not hot dogs. Our goal is to build a model that can accurately distinguish between the two classes.
One option would be to define a CNN architecture from scratch and train it on this small dataset directly. However, with such limited data, it‘s likely the model would overfit drastically and fail to generalize. Instead, we‘ll use transfer learning, starting with a pre-trained ResNet50 model. We‘ll use the feature extraction approach, freezing the pre-trained layers and only training a new binary classification layer on our hot dog data.
Here are the steps to perform transfer learning with ResNet50 in TensorFlow:
# Load pre-trained ResNet50 model
resnet_50 = tf.keras.applications.resnet50.ResNet50(include_top=False, weights=‘imagenet‘)
# Freeze base model layers
resnet_50.trainable = False
# Add new classification layers
inputs = keras.Input(shape=(224, 224, 3))
x = resnet_50(inputs)
x = keras.layers.GlobalAveragePooling2D()(x)
outputs = keras.layers.Dense(2, activation="softmax")(x)
# Define new model
model = keras.Model(inputs, outputs)
# Compile model
model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
# Train model on new data
model.fit(train_data, validation_data=test_data, epochs=10)Let‘s break this down step-by-step. First, we load the pre-trained ResNet50 model using tf.keras.applications. We exclude the top fully connected layer using include_top=False since we want to replace it with a new layer specific to our hot dog classification task. We also load the model with pre-trained ImageNet weights by specifying weights=‘imagenet‘.
Next, we freeze all the layers of the base ResNet model by setting resnet_50.trainable = False. This means that the original ImageNet weights of these layers will not be updated during training on our new hot dog data. The earlier layers have already learned general features like edges, textures, and shapes that will be useful for our task without modification. We‘re simply training a new final layer to adapt these features to hot dog classification.
Now we can add this new final layer to perform binary classification. We use a keras.Input layer to define the expected input shape (224 x 224 RGB images), pass this through the ResNet50 base model, then a GlobalAveragePooling2D layer to flatten the feature maps, and finally a Dense layer with 2 nodes and softmax activation to output class probabilities (hot dog vs. not hot dog). We use a softmax final layer since this is a multi-class problem – the 2 probabilities will always sum to 1. If we were doing binary classification with a single output node, we would use sigmoid activation instead.
We then define a new Model connecting the inputs to the outputs, compile it using the Adam optimizer and categorical cross-entropy loss, and train it on our hot dog datasets. After just 10 epochs, the model achieves over 80% validation accuracy!
This demonstrates the basic recipe for feature extraction with a pre-trained model in TensorFlow. To further improve performance, we could explore fine-tuning some of the deeper ResNet layers, experimenting with different model architectures, adding data augmentation, adjusting hyperparameters, and more. Even without these refinements, it‘s remarkable that with just 500 training examples, we can build an effective hot dog classifier by recycling knowledge learned from 1.2 million ImageNet examples.
Looking ahead, transfer learning will only grow more powerful and pervasive. Researchers are pushing the boundaries with techniques like few-shot and zero-shot learning to radically reduce data requirements [10]. At the same time, the diversity and scale of pre-trained models is rapidly expanding across modalities and domains – for example, Google‘s MultiModel architecture can simultaneously handle images, text and speech [11]. With tools like TensorFlow and Keras, anyone can tap into these vast reservoirs of pre-trained knowledge to solve data-scarce problems.
So whether you‘re a medical researcher classifying tumor biopsies, a conservationist trying to track endangered species, or an indie mobile developer building the next meme generator, you should absolutely consider transfer learning as a way to overcome limited data and computational resources. Thanks to the power of open-source libraries and pre-trained model zoos, state-of-the-art AI is now at your fingertips.
