Using PyTorch Lightning Flash and FiftyOne to Build Computer Vision Models, Fast | by Eric Hofesmann | Towards Data Science

Open-source tools have made significant advances in recent years to fill many of the same needs as end-to-end platform services. They can be incredibly useful for everything from model architecture development, to dataset curation, to model training and deployment. With enough digging, you can find an open-source tool that is able to support most portions of the data and model lifecycle. Tight integrations between tools are the best way to enable near-seamless workflows. This post delves into a new integration between the model prototyping and training framework, PyTorch Lightning Flash, and the dataset visualization and model analysis tool, FiftyOne.
Lightning Flash is a new framework built atop PyTorch Lighting and provides a collection of tasks for fast prototyping, baselining, fine-tuning, and solving business and scientific problems with deep learning. Even though Flash is easy to pick up no matter the amount of experience you have with deep learning, you can modify the existing tasks with Lightning and PyTorch to find the right level of abstraction for you. To speed things up even more, Flash code is scalable with built-in support for distributed training and inference on any hardware.
Flash makes it really easy to train your first model, but to continue improving it you will need to know how your model performs and how you can improve it. FiftyOne is an open-source tool for building high-quality datasets and computer vision models developed by Voxel51. It provides the building blocks for optimizing your dataset analysis pipeline, allowing you to get hands-on with your data, including visualizing complex labels, evaluating your models, exploring scenarios of interest, identifying failure modes, finding annotation mistakes, curating training datasets, and much more.
With Flash + FiftyOne, you can load datasets, train models, and analyze results for all of the following computer vision tasks:
Image Classification
Image Object Detection
Image Semantic Segmentation
Video Classification
Embeddings Visualization
The tight integration between Flash and FiftyOne allows you to execute an end-to-end workflow of loading a dataset, training a model on it, and visualizing/analyzing its predictions, all with just a few simple blocks of code
While it has always been easy to develop datasets with FiftyOne, the integration with PyTorch Lightning Flash now allows you to put those datasets to work by loading them into Flash and training tasks directly.
Flash provides the tools needed to grab a model for your task and start finetuning it on your data with as little code as possible and, most importantly, without needing to be an expert in the field.
Visualization of image and video data is always a challenge due to the complexity and size of modern datasets. FiftyOne was designed to provide a user-friendly view into your datasets and labels (including annotations and model predictions) which is now accessible by Flash models in one additional line of code.
To follow along with the examples in this post, you will need to install relevant packages. Primarily, you will need to install PyTorch Lightning Flash and FiftyOne.
pip install fiftyone lightning-flash
For the embeddings visualization workflow, you will also need to install the dimensionality reduction package, umap-learn:
pip install umap-learn
Most model development workflows with these tools follow the same general structure:
Load a dataset into FiftyOneCreate a Flash datamodule from the datasetFinetune a taskGenerate predictions from a modelAdd the predictions back to the dataset and visualize them
Load a dataset into FiftyOne
Create a Flash datamodule from the dataset
Finetune a task
Generate predictions from a model
Add the predictions back to the dataset and visualize them
This section shows a concrete example of using this integration between PyTorch Lightning Flash and FiftyOne to train and evaluate an image object detection model.
From here, you now have your predictions back in your dataset and can run an evaluation to generate confusion matrices, PR curves, and metrics like accuracy and mAP. In particular, you are able to identify and see individual true/false positive/negative results that allow you to understand where your model performs well and where it performs poorly. Improving your model based on common failure modes is the more surefire way to develop better models.
This workflow is unique in that it takes pretrained models and uses them to generate embedding vectors for every image in your dataset. You can then compute a visualization of these embeddings in a low dimensional space to find clusters of data. This functionality can lead to invaluable findings for the purposes of hard sample mining, data pre-annotation, sample recommendation for annotation, and much more.
With the notion of interactive plots, you are able to click on or circle regions of these embeddings and automatically update your session to view and tag the corresponding samples.
You can view similar workflows for other tasks like classification and segmentation here.
Over the years, impressive developments have come from the open-source community, especially in the field of machine learning. While individual tools can be great at solving specific problems, it is the tight integrations between tools that lead to powerful workflows. The new integration between PyTorch Lightning Flash and FiftyOne provides a new and easy way to develop datasets, train models, and analyze results.
(This post was co-authored by the PyTorch Lightning and Voxel51 teams)