We, at Lightly, are passionate engineers who want to make deep learning more efficient. That's why - together with our community - we want to popularize the use of self-supervised methods to understand and curate raw image data. Our solution can be applied before any data annotation step and the learned representations can be used to visualize and analyze datasets. This allows to select the best set of samples for model training through advanced filtering.
Lightly offers features like
You can find sample code for all the supported models here. We provide PyTorch, PyTorch Lightning, and PyTorch Lightning distributed examples for all models to kickstart your project.
Models:
Want to jump to the tutorials and see Lightly in action?
Tutorials for the Lightly Solution (Lightly Worker & API):
Community and partner projects:
Lightly requires Python 3.6+ but we recommend using Python 3.7+. We recommend installing Lightly in a Linux or OSX environment.
Lightly is compatible with PyTorch and PyTorch Lightning v2.0+!
Vision transformer based models require Torchvision v0.12+.
You can install Lightly and its dependencies from PyPI with:
pip3 install lightly
We strongly recommend that you install Lightly in a dedicated virtualenv, to avoid conflicting with your system packages.
If you only want to install the API client without torch and torchvision dependencies follow the docs on how to install the Lightly Python Client.
With Lightly, you can use the latest self-supervised learning methods in a modular way using the full power of PyTorch. Experiment with different backbones, models, and loss functions. The framework has been designed to be easy to use from the ground up. Find more examples in our docs.
import torch
import torchvision
from lightly import loss
from lightly import transforms
from lightly.data import LightlyDataset
from lightly.models.modules import heads
# Create a PyTorch module for the SimCLR model.
class SimCLR(torch.nn.Module):
def __init__(self, backbone):
super().__init__()
self.backbone = backbone
self.projection_head = heads.SimCLRProjectionHead(
input_dim=512, # Resnet18 features have 512 dimensions.
hidden_dim=512,
output_dim=128,
)
def forward(self, x):
features = self.backbone(x).flatten(start_dim=1)
z = self.projection_head(features)
return z
# Use a resnet backbone.
backbone = torchvision.models.resnet18()
# Ignore the classification head as we only want the features.
backbone.fc = torch.nn.Identity()
# Build the SimCLR model.
model = SimCLR(backbone)
# Prepare transform that creates multiple random views for every image.
transform = transforms.SimCLRTransform(input_size=32, cj_prob=0.5)
# Create a dataset from your image folder.
dataset = data.LightlyDataset(input_dir="./my/cute/cats/dataset/", transform=transform)
# Build a PyTorch dataloader.
dataloader = torch.utils.data.DataLoader(
dataset, # Pass the dataset to the dataloader.
batch_size=128, # A large batch size helps with the learning.
shuffle=True, # Shuffling is important!
)
# Lightly exposes building blocks such as loss functions.
criterion = loss.NTXentLoss(temperature=0.5)
# Get a PyTorch optimizer.
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, weight_decay=1e-6)
# Train the model.
for epoch in range(10):
for (view0, view1), targets, filenames in dataloader:
z0 = model(view0)
z1 = model(view1)
loss = criterion(z0, z1)
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(f"loss: {loss.item():.5f}")
You can easily use another model like SimSiam by swapping the model and the loss function.
# PyTorch module for the SimSiam model.
class SimSiam(torch.nn.Module):
def __init__(self, backbone):
super().__init__()
self.backbone = backbone
self.projection_head = heads.SimSiamProjectionHead(512, 512, 128)
self.prediction_head = heads.SimSiamPredictionHead(128, 64, 128)
def forward(self, x):
features = self.backbone(x).flatten(start_dim=1)
z = self.projection_head(features)
p = self.prediction_head(z)
z = z.detach()
return z, p
model = SimSiam(backbone)
# Use the SimSiam loss function.
criterion = loss.NegativeCosineSimilarity()
You can find a more complete example for SimSiam here.
Use PyTorch Lightning to train the model:
from pytorch_lightning import LightningModule, Trainer
class SimCLR(LightningModule):
def __init__(self):
super().__init__()
resnet = torchvision.models.resnet18()
resnet.fc = torch.nn.Identity()
self.backbone = resnet
self.projection_head = heads.SimCLRProjectionHead(512, 512, 128)
self.criterion = loss.NTXentLoss()
def forward(self, x):
features = self.backbone(x).flatten(start_dim=1)
z = self.projection_head(features)
return z
def training_step(self, batch, batch_index):
(view0, view1), _, _ = batch
z0 = self.forward(view0)
z1 = self.forward(view1)
loss = self.criterion(z0, z1)
return loss
def configure_optimizers(self):
optim = torch.optim.SGD(self.parameters(), lr=0.06)
return optim
model = SimCLR()
trainer = Trainer(max_epochs=10, devices=1, accelerator="gpu")
trainer.fit(model, dataloader)
See our docs for a full PyTorch Lightning example.
Or train the model on 4 GPUs:
# Use distributed version of loss functions.
criterion = loss.NTXentLoss(gather_distributed=True)
trainer = Trainer(
max_epochs=10,
devices=4,
accelerator="gpu",
strategy="ddp",
sync_batchnorm=True,
use_distributed_sampler=True, # or replace_sampler_ddp=True for PyTorch Lightning <2.0
)
trainer.fit(model, dataloader)
We provide multi-GPU training examples with distributed gather and synchronized BatchNorm. Have a look at our docs regarding distributed training.
Implemented models and their performance on various datasets. Hyperparameters are not tuned for maximum accuracy. For detailed results and more info about the benchmarks click here.
Note: Evaluation settings are based on these papers:
See the benchmarking scripts for details.
Model | Backbone | Batch Size | Epochs | Linear Top1 | Finetune Top1 | KNN Top1 | Tensorboard | Checkpoint |
---|---|---|---|---|---|---|---|---|
DINO | Res50 | 128 | 100 | 68.2 | 72.5 | 49.9 | link | link |
SimCLR | Res50 | 256 | 100 | 63.2 | N/A | 44.9 | link | link |
SwAV | Res50 | 256 | 100 | 67.2 | 75.4 | 49.5 | link | link |
Model | Backbone | Batch Size | Epochs | KNN Top1 |
---|---|---|---|---|
BarlowTwins | Res18 | 256 | 800 | 0.852 |
BYOL | Res18 | 256 | 800 | 0.887 |
DCL | Res18 | 256 | 800 | 0.861 |
DCLW | Res18 | 256 | 800 | 0.865 |
DINO | Res18 | 256 | 800 | 0.888 |
FastSiam | Res18 | 256 | 800 | 0.873 |
MAE | ViT-S | 256 | 800 | 0.610 |
MSN | ViT-S | 256 | 800 | 0.828 |
Moco | Res18 | 256 | 800 | 0.874 |
NNCLR | Res18 | 256 | 800 | 0.884 |
PMSN | ViT-S | 256 | 800 | 0.822 |
SimCLR | Res18 | 256 | 800 | 0.889 |
SimMIM | ViT-B32 | 256 | 800 | 0.343 |
SimSiam | Res18 | 256 | 800 | 0.872 |
SwaV | Res18 | 256 | 800 | 0.902 |
SwaVQueue | Res18 | 256 | 800 | 0.890 |
SMoG | Res18 | 256 | 800 | 0.788 |
TiCo | Res18 | 256 | 800 | 0.856 |
VICReg | Res18 | 256 | 800 | 0.845 |
VICRegL | Res18 | 256 | 800 | 0.778 |
Model | Backbone | Batch Size | Epochs | KNN Top1 |
---|---|---|---|---|
BarlowTwins | Res18 | 512 | 800 | 0.859 |
BYOL | Res18 | 512 | 800 | 0.910 |
DCL | Res18 | 512 | 800 | 0.874 |
DCLW | Res18 | 512 | 800 | 0.871 |
DINO | Res18 | 512 | 800 | 0.848 |
FastSiam | Res18 | 512 | 800 | 0.902 |
Moco | Res18 | 512 | 800 | 0.899 |
NNCLR | Res18 | 512 | 800 | 0.892 |
SimCLR | Res18 | 512 | 800 | 0.879 |
SimSiam | Res18 | 512 | 800 | 0.904 |
SwaV | Res18 | 512 | 800 | 0.884 |
SMoG | Res18 | 512 | 800 | 0.800 |
Below you can see a schematic overview of the different concepts in the package. The terms in bold are explained in more detail in our documentation.
Head to the documentation and see the things you can achieve with Lightly!
To install dev dependencies (for example to contribute to the framework) you can use the following command:
pip3 install -e ".[dev]"
For more information about how to contribute have a look here.
Unit tests are within the tests directory and we recommend running them using pytest. There are two test configurations available. By default, only a subset will be run:
make test-fast
To run all tests (including the slow ones) you can use the following command:
make test
To test a specific file or directory use:
pytest <path to file or directory>
To format code with black and isort run:
make format
Self-Supervised Learning:
Why should I care about self-supervised learning? Aren't pre-trained models from ImageNet much better for transfer learning?
How can I contribute?
Is this framework for free?
If this framework is free, how is the company behind Lightly making money?
If you want to cite the framework feel free to use this:
@article{susmelj2020lightly,
title={Lightly},
author={Igor Susmelj and Matthias Heller and Philipp Wirth and Jeremy Prescott and Malte Ebner et al.},
journal={GitHub. Note: https://github.com/lightly-ai/lightly},
year={2020}
}
Author: lightly-ai
Source Code: https://github.com/lightly-ai/lightly
License: MIT license
#machinelearning #python #deeplearning #computervision #pytorch