¶

Build better computer vision models faster with self-supervised pre-training

Why LightlyTrain?¶

LightlyTrain uses self-supervised learning (SSL) to train computer vision models on large datasets without labels. It provides simple Python, Command Line, and Docker interfaces to train models with popular SSL methods such as DINO and SimCLR.

The trained models are ideal starting points for fine-tuning on downstream tasks such as image classification, object detection, and segmentation or for generating image embeddings. Models trained with LightlyTrain result in improved performance, faster convergence, and better generalization compared to models trained without SSL. Image embeddings created with LightlyTrain capture more relevant information than their supervised counterparts and seamlessly extend to new classes due to the unsupervised nature of SSL.

Lightly is the expert in SSL for computer vision and developed LightlyTrain to simplify model training for any task and dataset.

How It Works¶

Train a Model with SSL¶

import lightly_train

ligthly_train.train(
    out="out/my_experiment",            # Output directory
    data="my_data_dir",                 # Directory with images
    model="torchvision/resnet50",       # Model to train
    method="dino",                      # Self-supervised learning method
)

After training is complete, you can either export the model for fine-tuning or use the model to generate image embeddings.

Export the Model for Fine-Tuning¶

import lightly_train

lightly_train.export(
    out="my_exported_model.pth",                            # Exported model
    checkpoint="out/my_experiment/checkpoints/last.ckpt",   # LightlyTrain checkpoint
    part="model",                                           # Model part to export
    format="torch_state_dict",                              # Export format
)

You can now fine-tune the exported model with your favorite library.

Generate Image Embeddings¶

import lightly_train

lightly_train.embed(
    out="my_embeddings.pth",                                # Exported embeddings
    checkpoint="out/my_experiment/checkpoints/last.ckpt",   # LightlyTrain checkpoint
    data="my_data_dir",                                     # Directory with images
    format="torch",                                         # Embedding format
)

You can now use the generated embeddings for clustering, retrieval, or visualization tasks.

The quick start guide shows in more detail how to install and use LightlyTrain.

Features¶

Train models on any image data without labels
Train models from popular libraries such as torchvision, TIMM, Ultralytics, and SuperGradients
Train custom models with ease
No self-supervised learning expertise required
Automatic SSL method selection (coming soon!)
Python, Command Line, and Docker support
Multi-GPU and multi-node support (coming soon!)
Export models for fine-tuning or inference
Generate and export image embeddings
Monitor training progress with TensorBoard, Weights & Biases, and more