vturrisi/solo-learn

solo-learn: a library of self-supervised methods for visual representation learning powered by Pytorch Lightning

What it solves

solo-learn provides a unified, comparable environment for implementing and evaluating state-of-the-art self-supervised methods for unsupervised visual representation learning. It eliminates the need to implement these methods from scratch and includes various training tricks to improve performance and efficiency.

How it works

Built on PyTorch Lightning, the library implements a wide array of self-supervised learning (SSL) algorithms (such as SimCLR, BYOL, MoCo, and MAE) and integrates them with various neural network backbones (like ResNet, ViT, and ConvNeXt). It uses a flexible configuration system via OmegaConf/Hydra and supports high-speed data processing through Nvidia Dali.

Who it’s for

It is designed for researchers and developers working on computer vision who want to benchmark different self-supervised learning methods, implement new SSL techniques, or pretrain visual backbones without labeled data.

Highlights

  • Extensive Method Library: Supports a vast range of SSL methods including All4One, Barlow Twins, DINO, and SwAV.
  • Diverse Backbones: Compatible with ResNet, WideResNet, ViT, Swin, PoolFormer, and ConvNeXt.
  • Advanced Evaluation: Offers online linear evaluation (via stop-gradient), offline linear evaluation, K-NN evaluation, and UMAP feature space visualization.
  • Performance Optimizations: Includes channel last conversion and Nvidia Dali integration for significantly faster data processing.
  • Training Utilities: Integrated with WandB for cloud logging and provides custom model checkpointing.

Related

  • Project
    lightly-ai/lightly

    A computer vision framework for self-supervised learning that provides modular building blocks and implementations of numerous state-of-the-art SSL models.

  • Project
    Lightning-AI/pytorch-lightning

    A deep learning framework that abstracts away the engineering boilerplate of PyTorch, enabling researchers to pretrain and finetune AI models at scale across various hardware.

  • Project
    hmorimitsu/ptlflow

    A unified PyTorch Lightning framework for training and evaluating a wide collection of state-of-the-art optical flow estimation models.

  • Project
    pytorch-tabular/pytorch_tabular

    A high-level PyTorch framework that provides a unified interface for applying various deep learning architectures to tabular data.

  • Project
    yihaosun1124/OfflineRL-Kit

    A PyTorch library for offline reinforcement learning that provides a standardized framework and implementations of state-of-the-art model-free and model-based algorithms.