PyTorch Lightning Bolts — From Linear, Logistic Regression on TPUs to pre-trained GANs
By: Pytorch Lightning Team
Tags: Kaggle, Data Science, Machine Learning, Pytorch Lightning, Pytorch

Lightning Bolts Use Cases

Tested Baselines

Bolts has rigorously tested and benchmarked baselines. From VAEs to GANs to GPT to self-supervised models — you don’t have to spend months implementing the baselines to try new ideas. Instead, subclass one of ours and try your idea!

Predicting on your data

You can use Bolts pre-trained weights on most of the standard datasets. This is useful when you don’t have enough data, time, or money to do your own training. It can also be used as a benchmark to improve or test your own model.

For example, you could use a pre-trained VAE to generate features for an image dataset, and compare it to CPC.

Subclass and train

Bolts modules are written to be easily extended for research. For example, you can subclass SimCLR and make changes to the NT-Xent loss.

Mix and Match

The beauty of Bolts is that it’s easy to plug and play with your Lightning modules or any PyTorch data set. Mix and match data, modules, and components as you please!

model = ImageGPT(datamodule=FashionMNISTDataModule(PATH))

Or pass in any dataset of your choice

model = ImageGPT()

Trainer().fit(model,

train_dataloader=DataLoader(…),

val_dataloader=DataLoader(…)

)

And train on any hardware accelerator, using all of Lightning trainer flags!

import pytorch_lightning as pl model = ImageGPT(datamodule=FashionMNISTDataModule(PATH)) # cpus

pl.Trainer().fit(model) # gpus

pl.Trainer(gpus=8).fit(model) # tpus

pl.Trainer(tpu_cores=8).fit(model)

For more details, read the docs.

Logistic regressions on multiple-GPUs and TPUs

Lightning Bolts includes a collection of non-deep learning algorithms that can train on multiple GPUs and TPUs.

Here’s an example running logistic regression on Imagenet in 2 GPUs with 16-bit precision.

This is the summary for the above setup:

1,281,167 images (in Imagenet)

150,528 input features (224x224x3 pixels)

512 batch size (256 per GPU)

16-bit precision

2 V100 GPUs

Self Supervised Learning Models

Contrastive learning

Bolts houses a collection of many of the current state-of-the-art self-supervised algorithms.

In the latest paper, “A Framework For Contrastive Self-Supervised Learning And Designing A New Approach” by William Falcon and Kyunghyun Cho, explore and implement many of the latest self-supervised learning algorithms. Those implementations are available in Bolts under the self-supervised learning suite.

You can use Self-supervised learning algorithms to extract image features, to train unlabeled data or mix and match parts of the models to create your own new method:

from pl_bolts.models.self_supervised import CPCV2

from pl_bolts.losses.self_supervised_learning import FeatureMapContrastiveTask



amdim_task = FeatureMapContrastiveTask(comparisons='01, 11, 02', bidirectional=True)

model = CPCV2(contrastive_task=amdim_task)

You can find bolts implementation for

Contribute to Bolts!

Bolts models and components are built-by the Lightning community. The lightning team guarantees that contributions are:

Rigorously Tested (on CPUs, GPUs, TPUs). Rigorously Documented. Standardized via PyTorch Lightning. Optimized for reproducibility. Checked for correctness.

Want to get your implementation tested on CPUs, GPUs, TPUs, and mixed-precision and help us grow? Consider contributing your model to Bolts (you can even do it from your own repo) to make it available for the Lightning community!

You can find a list of model suggestions in our GitHub issues. Your models don’t have to be state of the art to get added to bolts, just well designed and tested. If you have any ideas, start a discussion on our slack channel.