Introducing Lightning Flash — From Deep Learning Baseline To Research in a Flash

Flash is a collection of tasks for fast prototyping, baselining and fine-tuning scalable Deep Learning models, built on PyTorch Lightning.

Whether you are new to deep learning, or an experienced researcher, Flash offers a seamless experience from baseline experiments to state-of-the-art research. It allows you to build models without being overwhelmed by all the details, and then seamlessly override and experiment with Lightning for full flexibility. Continue reading to learn how to use Flash tasks to get state-of-the-art results in a flash.

Why Flash?

PyTorchLightning/lightning-flash

1. The power of lightning, without the prerequisites

Over the past year, PyTorch Lightning has received an enthusiastic response from the community for decoupling research from boilerplate code, enabling seamless distributed training, logging, and reproducibility of deep learning research code. Lightning is used by top research labs and AI companies to simplify the training of PyTorch models around the world.

However, getting started with Deep Learning projects can still be quite intimidating. Even for an expert, it can take quite some time to get a baseline model running on a new dataset or out of domain task.

With Flash, you can create your own image or text classifier in a few lines of code, no math, fancy nn.Modules or research experience required!

2. Quickest way to a deep learning baseline

The typical workflow for trying out a new Deep Learning task follows the three-step cycle below:

This cycle is a major bottleneck to Deep Learning adoption

This cycle is a major bottleneck to Deep Learning adoption. Most example repositories do not scale to production training and inference. They are often not maintained for more than a few weeks or at best months at a time. Data loading is typically hardcoded to different benchmark datasets which leads to numerous incompatible standards for representing Tasks. Even worse, it is usually near impossible to modify or build upon the code until you master how it works.

We created Flash to answer our user’s need for a super quick way to get a Lightning baseline to benchmark their experiments against, using proven backbones for common data patterns. Flash replaces this cycle enabling our users to focus on science, not infrastructure or ops. Flash makes the power of Lightning more accessible to data scientists, developers and Kagglers, and makes baselining trivial for more experienced researchers.

3. Scalability

Flash is the first high-level framework to provide seamless support for distributed training and inference of Deep Learning models. With the power of Lightning, you can train or fine-tune your flash tasks on any hardware: CPUs, GPUs, or TPUs without any code changes.

4. Flexibility

Consistent with PyTorch Lightning’s goal of getting rid of the boilerplate, Flash’s vision is to make it easy to train, inference and finetune models with Lightning as quickly and as flexibly as possible. With Flash, as you learn more, you can override your Task code with Lightning and PyTorch to find the right level of abstraction for your skillset.

How Flash Works

Flash is comprised of a collection of Tasks. Flash Tasks are laser-focused objects designed to solve common problems, using state-of-the-art methods. Flash Tasks are designed to make Inference, Finetuning, and Training seamless.

We currently support Image classification, Image embedding, tabular classification, text classification, summarization, translation, but many other tasks are already in the works!

Flash tasks contain all the relevant information to solve the task at hand- the number of class labels you want to predict, the number of columns in your dataset, as well as details on the model architecture used such as loss function, optimizers, etc. You can choose which architecture to use for a variety of implementation by changing the backbone , overriding the loss function or optimizers, etc.

Using Flash

Say you have a dataset from a Kaggle challenge, a research project, or an app you are building. Flash tasks enable you to infer, finetune and train models on your data.

Inference

Inference is the process of generating predictions from trained models. All Flash tasks have pre-trained weights you can use out-of-the-box.

To use Flash tasks for inference, you simply:

1. Initialize your task with pre-trained weights using a checkpoint (a checkpoint is simply a file that captures the exact value of all parameters used by a model). Any local file or remote URL works.
2. Pass in the data to flash.core.model.Task.predict().

For example, the ImageClassifier was trained on ImageNet dataset, to classify 10 classes of animals. We can use the pre-trained weights to predict if an image and an ant or a bee, using the Hymenoptera data set. In this example, we load the pre-trained model from a checkpoint and use the Hymenoptera data set to make predictions.

Finetune

Finetuning (or transfer-learning) is the process of tweaking a model trained on a large dataset, to your particular (likely much smaller) dataset. To use a Task for finetuning:

1. Download and set up your own data (DataLoader or LightningModule work).
2. Initialize your task.
3. Initialize a flash.core.trainer.Trainer (or a Lightning Trainer). The trainer allows you to train or finetune your task.
4. Call Trainer.finetune() with your task and data set.
5. Use your fine-tuned model for predictions.

Here’s an example of how simple fine-tuning is with Flash:

Train from scratch

When you have enough data, it is better to train from scratch instead of finetuning.

Steps for training are similar to finetuning:

1. Download and set up your own data (DataLoader or LightningModule work).
2. Init your task.
3. Init a flash.core.trainer.Trainer (or a Lightning Trainer).
4. Call flash.core.trainer.Trainer.fit() with your task and data set.
5. Use your fine-tuned model for predictions

Learn More

Now that you have the tools to get started with Deep Learning, we cannot wait for you to show us what you can build.

We are working tirelessly on adding more tasks to Flash, so if you have any must-have tasks comment below or reach out to us on twitter @pytorchlightnin or in our Slack channel.

— —
Note this article was edited on February 4th, 2021 to clarify the vision of Flash.