/byol

PyTorch implementation of BYOL: a fantastically simple method for self-supervised image representation learning with SOTA performance.

Primary LanguagePythonMIT LicenseMIT

BYOL: Bootstrap Your Own Latent

PyTorch implementation of BYOL: a fantastically simple method for self-supervised image representation learning with SOTA performance. Strongly influenced and inspired by this Github repo, but with a few notable differences:

  1. Enables multi-GPU training in PyTorch Lightning.
  2. (Optionally) Automatically trains a linear classifier, and logs its accuracy after each epoch.
  3. All functions and classes are fully type annotated for better usability/hackability with Python>=3.6.

TO DO

  • Enable mixed-precision training in PyTorch Lightning. kornia.augmentation.RandomResizedCrop currently doesn't support this. I'll need to ensure that our implementation is sufficiently performant, so it doesn't inadvertently slow down training.