Lightning Flash is a high-level deep learning framework for fast prototyping, baselining, fine-tuning, and solving deep learning problems. Flash makes complex AI recipes for over 15 tasks across 7 data domains accessible to all. It is built for beginners with a simple API that requires very little deep learning background, and for data scientists, Kagglers, applied ML practitioners, and deep learning researchers that want a quick way to get a deep learning baseline with advanced features PyTorch Lightning offers.
pip install lightning-flash
Using Lightning Flash in 3 Steps!¶
1. Load your Data¶
All data loading in Flash is performed via a
from_* classmethod of a
DataModule to use and which
from_* methods are available depends on the task you want to perform.
For example, for image segmentation where your data is stored in folders, you would use the
SemanticSegmentationData’s from_folders method:
from flash.image import SemanticSegmentationData dm = SemanticSegmentationData.from_folders( train_folder="data/CameraRGB", train_target_folder="data/CameraSeg", val_split=0.1, image_size=(256, 256), num_classes=21, )
2. Configure your Model¶
Our tasks come loaded with pre-trained backbones and (where applicable) heads. You can view the available backbones to use with your task using available_backbones. Once you’ve chosen, create the model:
from flash.image import SemanticSegmentation print(SemanticSegmentation.available_heads()) # ['deeplabv3', 'deeplabv3plus', 'fpn', ..., 'unetplusplus'] print(SemanticSegmentation.available_backbones("fpn")) # ['densenet121', ..., 'xception'] # + 113 models print(SemanticSegmentation.available_pretrained_weights("efficientnet-b0")) # ['imagenet', 'advprop'] model = SemanticSegmentation(head="fpn", backbone="efficientnet-b0", pretrained="advprop", num_classes=dm.num_classes)
from flash import Trainer trainer = Trainer(max_epochs=3) trainer.finetune(model, datamodule=datamodule, strategy="freeze") trainer.save_checkpoint("semantic_segmentation_model.pt")
To learn more about Lightning Flash, please refer to the Lightning Flash documentation.