Setting A Strong Deep Learning Baseline In Minutes With PyTorch

Iterate your way from baseline to custom models to ship products faster or to publish your research faster.

William Falcon
Towards Data Science

--

With Flash build PyTorch baselines in minutes

Whether you’re a data scientist, research engineer, AI researcher, or machine learning engineer, baselines are non-negotiable. Don’t build a fancy GAN or try a complex idea before setting up a good foundation.

In this tutorial, we’ll use Flash to build two PyTorch baselines in minutes. After that, we’ll iterate on that baseline using Lightning to get to a custom implementation tailored to your particular circumstances and squeeze more performance from your models.

Baseline Vs End-Model

To get value fast, iterate from baseline to end-model in phases.

Most of the time, we want to get to the best solution as fast as possible with deep learning. But that might take months of research and paper implementation. Along that journey, you can use a few tricks to get you to that optimal model iteratively, by starting with a strong baseline.

Remember, the goal of a baseline is to get your idea in front of users or get a benchmark on the dataset you’re working on. Once you’ve established your baselines, you can move on to all the complex deep learning engineering you need to do knowing you’re starting off from a strong base.

Baseline 1: Ants vs Bees

To illustrate the key ideas, let's say you work for a home insurance agency and you developed a camera that owners install to track the types of bugs in their house. Your insurance is based on the number of Ants and Bees that the camera captures. As part of this system, you decide to build a classifier to figure out what is an Ant or a Bee.

First, let’s install lightning-flash.

pip install lightning-flash

Now what we need to do is to find a strong pretrained model and finetune on your data. Here we use a model pretrained on Imagenet and adjust the weights using the “freeze_unfreeze” scheme.

In just a few lines you’ve managed to get a baseline done for your work. You run it and see that it gets you 80% prediction accuracy. If that’s good enough for you, you can deploy this model within the same day.

Baseline 2: Fancier Baseline

Another easy win is to change the backbone of your baseline to use the latest self-supervised model from Facebook AI (Swav). With Flash this is simple:

Just change the backbone from ‘resnet34’ to ‘swav-imagenet’.

ImageClassifier(backbone='swav-imagenet')

Now if you’re lucky, your model gets another 5% accuracy. Within the same day, you’ve iterated on your work twice!

The End-Model

But now is where the fun part and actual work begins.

Here you’ll brush off all those dusty math books, find a list of deep learning papers and start experimenting with the latest and greatest in deep learning.

Let’s say that it turns out you have some other data you’d like to integrate into your model and you want to use a fancy scheduling technique that uses a fancy GAN regularizer term to do something super unique that could only work for you given your particular dataset.

Since Flash is a companion to Lightning, you can simply rewrite your baseline into a LightningModule and do all the fancy tricks you want.

Notice that you were able to go from baseline to super customized implementation without introducing a ton of boilerplate. You also maintain full control and customization all the way from a fast high-level baseline to a super fancy custom implementation.

Bonus, Deploy!

As a bonus, all LightningModules are easily exportable via torchscript or onnx. You can use these to deploy high performance models.

Summary

This article showed you how to build deep learning models iteratively to deliver value as fast as possible. Using the family of Lightning frameworks you were able to:

  • Deliver baseline v1 in minutes
  • Iterate to baseline v2 in a few more minutes
  • Iterate to the end-model (custom implementation) within hours/days.

Build and ship/publish fast!

Visit the Flash GitHub repo for more examples of rapid baselining!

--

--

⚡️PyTorch Lightning Creator • PhD Student, AI (NYU, Facebook AI research).