¶

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

lightly_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
)

This will pretrain a Torchvision ResNet-50 model using images from my_data_dir and the DINO self-supervised learning method. All training logs, model exports, and checkpoints are saved to the output directory at out/my_experiment.

The final model is exported to out/my_experiment/exported_models/exported_last.pt in the default format of the used library. It can directly be used for fine-tuning. Follow the example to learn how to fine-tune a model.

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

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
Built for high performance including Multi-GPU and multi-node support
Export models for fine-tuning or inference
Generate and export image embeddings
Monitor training progress with TensorBoard, Weights & Biases, and more