Your PyTorch AI Factory - Flash enables you to easily configure and run complex AI recipes.
Your PyTorch AI Factory
Installation • Flash in 3 Steps • Docs • Contribute • Community • Website • License
From PyPI:
pip install lightning-flash
See our installation guide for more options.
All data loading in Flash is performed via a from_* classmethod on a DataModule.
To decide which DataModule to use and which from_* methods are available, it depends on the task you want to perform.
For example, for image segmentation where your data is stored in folders, you would use the from_folders method of the SemanticSegmentationData class:
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,
)
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 one, 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")
Serve in just 2 lines:
from flash.image import SemanticSegmentation
model = SemanticSegmentation.load_from_checkpoint("semantic_segmentation_model.pt")
model.serve()
or make predictions from raw data directly.
from flash import Trainer
trainer = Trainer(strategy='ddp', accelerator="gpu", gpus=2)
dm = SemanticSegmentationData.from_folders(predict_folder="data/CameraRGB")
predictions = trainer.predict(model, dm)
Training strategies are PyTorch SOTA Training Recipes which can be utilized with a given task.
Check out this example where the ImageClassifier supports 4 Meta Learning Algorithms from Learn2Learn.
This is particularly useful if you use this model in production and want to make sure the model adapts quickly to its new environment with minimal labelled data.
from flash.image import ImageClassifier
model = ImageClassifier(
backbone="resnet18",
optimizer=torch.optim.Adam,
optimizer_kwargs={"lr": 0.001},
training_strategy="prototypicalnetworks",
training_strategy_kwargs={
"epoch_length": 10 * 16,
"meta_batch_size": 4,
"num_tasks": 200,
"test_num_tasks": 2000,
"ways": datamodule.num_classes,
"shots": 1,
"test_ways": 5,
"test_shots": 1,
"test_queries": 15,
},
)
In detail, the following methods are currently implemented:
With Flash, swapping among 40+ optimizers and 15+ schedulers recipes are simple. Find the list of available optimizers, schedulers as follows:
from flash.image import ImageClassifier
ImageClassifier.available_optimizers()
# ['A2GradExp', ..., 'Yogi']
ImageClassifier.available_schedulers()
# ['CosineAnnealingLR', 'CosineAnnealingWarmRestarts', ..., 'polynomial_decay_schedule_with_warmup']
Once you've chosen, create the model:
#### The optimizer of choice can be passed as
from flash.image import ImageClassifier
# - String value
model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer="Adam", lr_scheduler=None)
# - Callable
model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer=functools.partial(torch.optim.Adadelta, eps=0.5), lr_scheduler=None)
# - Tuple[string, dict]: (The dict takes in the optimizer kwargs)
model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer=("Adadelta", {"epa": 0.5}), lr_scheduler=None)
#### The scheduler of choice can be passed as a
# - String value
model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer="Adam", lr_scheduler="constant_schedule")
# - Callable
model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer="Adam", lr_scheduler=functools.partial(CyclicLR, step_size_up=1500, mode='exp_range', gamma=0.5))
# - Tuple[string, dict]: (The dict takes in the scheduler kwargs)
model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer="Adam", lr_scheduler=("StepLR", {"step_size": 10}))
You can also register you own custom scheduler recipes beforeahand and use them shown as above:
from flash.image import ImageClassifier
@ImageClassifier.lr_schedulers_registry
def my_steplr_recipe(optimizer):
return torch.optim.lr_scheduler.StepLR(optimizer, step_size=10)
model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer="Adam", lr_scheduler="my_steplr_recipe")
Flash includes some simple augmentations for each task by default, however, you will often want to override these and control your own augmentation recipe.
To this end, Flash supports custom transformations with the InputTransform.
The InputTransform is like a callback for transforms, with hooks that can be used to apply transforms to samples or batches, on and off the device / accelerator.
In addition, hooks can be specialized to apply transforms only to the input or target.
With these hooks, complex transforms like MixUp can be implemented with ease.
Here's an example (with an albumentations transform thrown in too!):
import torch
import numpy as np
import albumentations
from flash import InputTransform
from flash.image import ImageClassificationData
from flash.image.classification.input_transform import AlbumentationsAdapter
def mixup(batch, alpha=1.0):
images = batch["input"]
targets = batch["target"].float().unsqueeze(1)
lam = np.random.beta(alpha, alpha)
perm = torch.randperm(images.size(0))
batch["input"] = images * lam + images[perm] * (1 - lam)
batch["target"] = targets * lam + targets[perm] * (1 - lam)
return batch
class MixUpInputTransform(InputTransform):
def train_input_per_sample_transform(self):
return AlbumentationsAdapter(albumentations.HorizontalFlip(p=0.5))
# This will be applied after transferring the batch to the device!
def train_per_batch_transform_on_device(self):
return mixup
datamodule = ImageClassificationData.from_folders(
train_folder="data/train",
transform=MixUpInputTransform,
batch_size=2,
)
Flash Zero is a zero-code machine learning platform built
directly into lightning-flash
using the Lightning CLI.
To get started and view the available tasks, run:
flash --help
For example, to train an image classifier for 10 epochs with a resnet50 backbone on 2 GPUs using your own data, you can do:
flash image_classification --trainer.max_epochs 10 --trainer.gpus 2 --model.backbone resnet50 from_folders --train_folder {PATH_TO_DATA}
The lightning + Flash team is hard at work building more tasks for common deep-learning use cases. But we're looking for incredible contributors like you to submit new tasks!
Join our Slack and/or read our CONTRIBUTING guidelines to get help becoming a contributor!
Note: Flash is currently being tested on real-world use cases and is in active development. Please open an issue if you find anything that isn't working as expected.
Flash is maintained by our core contributors.
For help or questions, join our huge community on Slack!
We’re excited to continue the strong legacy of opensource software and have been inspired over the years by Caffe, Theano, Keras, PyTorch, torchbearer, and fast.ai. When/if additional papers are written about this, we’ll be happy to cite these frameworks and the corresponding authors.
Flash leverages models from many different frameworks in order to cover such a wide range of domains and tasks. The full list of providers can be found in our documentation.
Please observe the Apache 2.0 license that is listed in this repository.
FAQs
Your PyTorch AI Factory - Flash enables you to easily configure and run complex AI recipes.
We found that lightning-flash demonstrated a healthy version release cadence and project activity because the last version was released less than a year ago. It has 2 open source maintainers collaborating on the project.
Did you know?
Socket for GitHub automatically highlights issues in each pull request and monitors the health of all your open source dependencies. Discover the contents of your packages and block harmful activity before you install or update your dependencies.
Security News
Research
The Socket Research Team breaks down a malicious wrapper package that uses obfuscation to harvest credentials and exfiltrate sensitive data.
Research
Security News
Attackers used a malicious npm package typosquatting a popular ESLint plugin to steal sensitive data, execute commands, and exploit developer systems.
Security News
The Ultralytics' PyPI Package was compromised four times in one weekend through GitHub Actions cache poisoning and failure to rotate previously compromised API tokens.