Transfer Learning Using PyTorch Lightning ⚡️

In this colab, we will extend the pipeline here to perform transfer learning with PyTorch Lightning.

Transfer Learning is a technique where the knowledge learned while training a model for "task" A and can be used for "task" B. Here A and B can be the same deep learning tasks but on a different dataset.

Setting up PyTorch Lightning and W&B

For this tutorial, we need PyTorch Lightning and Weights and Biases.

You're gonna need these imports.

Now you'll need to login to you wandb account.

The Dataset 💿

We will be using the StanfordCars dataset to train our image classifier. It contains 16,185 images of 196 classes of cars. The data is split into 8,144 training images and 8,041 testing images, where each class has been split roughly in a 50-50 split. Classes are typically at the level of Make, Model, Year, e.g. 2012 Tesla Model S or 2012 BMW M3 coupe.

LightingModule - Define the System

Let us look at the model definition to see how transfer learning can be used with PyTorch Lightning. In the LitModel class, we can use the pre-trained model provided by Torchvision as a feature extractor for our classification model. Here we are using ResNet-18. A list of pre-trained models provided by PyTorch Lightning can be found here.

• When pretrained=True, we use the pre-trained weights; otherwise, the weights are initialized randomly.
• If .eval() is used, then the layers are frozen.
• A single Linear layer is used as the output layer. We can have multiple layers stacked over the feature_extractor.

Setting the transfer argument to True will enable transfer learning.

Train your Model 🏋️‍♂️

To train the model, we instantiate the StanfordCarsDataModule and the LitModel along with the PyTorch Lightning Trainer. To the Trainer, we will pass the WandbLogger as the logger to use W&B to track the metrics during model training!

We are good to go! Let's train our model!

Now that the model is trained, let's see how it performs on the test set

Let's close our W&B run, so we call wandb.finish().

The workspace generated to compare training the model from scratch vs using transfer learning is available here. The conclusions that can be drawn from this are explained in detail in this report.

Conclusion

I will encourage you to play with the code and train an image classifier with a dataset of your choice from scratch and using transfer learning.

To learn more about transfer learning check out these resources:

• Gotchas of transfer learning for image classification by Sayak Paul.
• Transfer Learning with Keras and Deep Learning by PyImageSearch.
• Transfer Learning - Machine Learning's Next Frontier by Sebastian Ruder.