byol-pytorch: Wrapping an Existing Backbone for Self-Supervised Pretraining
Usable Implementation of "Bootstrap Your Own Latent" self-supervised learning, from Deepmind, in Pytorch
At a glance
- What is it?
- A PyTorch wrapper that turns any image-based network into a BYOL learner with two keyword arguments, plus a distributed trainer built on Hugging Face Accelerate. The trade-off is that you inherit the paper's hyperparameters and the augmentation pipeline unless you replace them.
- Who is it for?
- Adopt byol-pytorch if you already have a working PyTorch classifier or backbone and want to pretrain it on unlabelled images without writing the BYOL loss, projection head and momentum update yourself. Do not adopt it if your downstream task is segmentation, where the README points at the separate pixel-level-contrastive-learning repository, or if you need the augmentation policy expressed as data rather than as an nn.Module.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 141 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The labelling bill this wrapper is aimed at
Supervised image training needs a label per image. BYOL's premise, as the README puts it, is to avoid "having to designate negative pairs", which is what contrastive methods such as SimCLR require. The repository's own framing is direct: it offers "a module that one can easily wrap any image-based neural network (residual network, discriminator, policy network) to immediately start benefitting from unlabelled image data." The intended user is someone who already has a PyTorch model and a pile of unlabelled images, and who wants the backbone to come out better on a later supervised task without rebuilding the training loop from scratch. The README closes the pitch with "Now go save your organization from having to pay for labels", which sets the audience: teams with image data and a labelling budget they would rather not spend.
Two encoders, one loss, and a moving average you have to call
BYOL trains an online encoder against a target encoder. The target is not updated by gradients. In the default configuration it is an exponential moving average of the online encoder, and the caller advances it by calling learner.update_moving_average() after each optimizer step. Forgetting that call is the most likely silent failure in a hand-written loop, because training still runs and the loss still moves. The wrapper exposes projection_size, projection_hidden_size and moving_average_decay as keyword arguments, with the README stating these are "already set to what the paper has found optimal" and that moving_average_decay defaults to 0.99. The online network is your model plus a projection MLP plus a prediction MLP; the target side is the same architecture without the predictor. The README also offers the SimSiam variant from Kaiming He's paper by setting use_momentum = False, in which case the README states you "will no longer need to invoke update_moving_average". That single flag changes the training loop's shape, not just a hyperparameter.
Augmentation is an nn.Module, not a config file
By default the library applies the SimCLR augmentations, which the README notes are also the ones used in the BYOL paper. If you want your own, you pass augment_fn, and the README's example is a kornia.augmentation.RandomHorizontalFlip() inside an nn.Sequential. There is a second slot, augment_fn2, for the asymmetric view. The README observes that in the paper "one of the augmentations have a higher gaussian blur probability than the other", and its example pairs a plain flip pipeline with a flip plus kornia.filters.GaussianBlur2d((3, 3), (1.5, 1.5)). This is a meaningful design choice: the augmentation policy lives in Python objects, so it is versioned with your training script rather than with a YAML file. That is convenient for research and awkward for teams that expect to sweep augmentation strength without touching code.
Wiring it into a training loop, and the distributed path
Installation is a single command, pip install byol-pytorch. The minimal loop from the README constructs BYOL(resnet, image_size=256, hidden_layer='avgpool'), wraps learner.parameters() in torch.optim.Adam at lr=3e-4, and after loss.backward() and opt.step() calls learner.update_moving_average(). The improved backbone is then saved with torch.save(resnet.state_dict(), './improved-net.pt'), which is worth noting: the checkpoint is the wrapped network's state dict, not the learner's, so the projection and prediction heads are discarded on save. hidden_layer accepts a name or an index, and the README's custom-augmentation example uses hidden_layer = -2, so negative indexing into the module list is supported. For multi-GPU work the repository ships BYOLTrainer, which takes a Dataset plus image_size, hidden_layer, learning_rate, num_train_steps, batch_size and checkpoint_every, with the README stating checkpoints land in a ./checkpoints folder. Setup is accelerate config, launch is accelerate launch ./train.py. The README's own example uses MockDataset(256, 10000), so the trainer is demonstrated against a synthetic dataset rather than a real one.
Where the wrapper gets in the way
The hidden_layer argument is the sharpest edge. You must name a module whose output is the representation you want to pretrain, and the README's examples use 'avgpool' for torchvision ResNet models. If the name does not resolve, the failure surfaces at construction or at the first forward pass, and the library cannot tell you which layer you should have chosen. The second constraint is the augmentation default. SimCLR-style pipelines are built around large batches of crops; the README's own distributed example uses batch_size = 16, and nothing in the material says how that interacts with the default augmentation strength. Third, this is an image library. The README's framing is explicitly "any image-based neural network", and the augmentation examples are all spatial image transforms. There is no stated support for text, audio or tabular inputs. Finally, the README carries three updates citing later papers, including one that replaced batch norm with group norm plus weight standardization. Those are pointers to the literature, not features of this package, and the README does not claim the library implements them.
The SimSiam switch and what it costs you
Setting use_momentum = False is the clearest alternative approach inside the library itself. The README attributes the idea to a paper from Kaiming He and describes the change plainly: BYOL "does not even need the target encoder to be an exponential moving average of the online encoder". Turning the flag off removes the update_moving_average() call and the momentum decay from the picture entirely, which simplifies the loop and removes one class of silent bug. What you give up is the mechanism the original BYOL paper is built around. The README does not state whether the two variants reach comparable accuracy on any given dataset, and it does not present a comparison. Choosing between them is therefore an empirical question you have to answer on your own data, not one the documentation settles. Outside the library, the README points to lucidrains/pixel-level-contrastive-learning for segmentation, which extends the same idea to pixel-level learning. That is a different repository with a different objective, not a mode of this one.
Maintenance, licence and what the release history shows
The licence is MIT, which permits commercial use and modification provided the copyright notice and permission notice are retained; that is a summary of the identifier, not legal advice, and you should read the LICENSE file in the repository before relying on it. On maintenance, the release list shows 0.8.0 in November 2023, 0.8.1 in May 2024 and 0.8.2 in July 2024, with the repository's last push recorded in April 2026 and no archived flag. The gap between the last tagged release and the last push is worth noting if you depend on tagged versions rather than the master branch. Upgrade cost is bounded by the surface area: the public entry points shown in the README are the BYOL class, the BYOLTrainer class and MockDataset, and the keyword arguments listed are image_size, hidden_layer, projection_size, projection_hidden_size, moving_average_decay, use_momentum, augment_fn, augment_fn2 and return_embedding. A minor version bump that changes the default augmentation pipeline would change training results without changing your code, so pinning the version is the practical precaution.
Editorial conclusion
Adopt byol-pytorch if you already have a working PyTorch classifier or backbone and want to pretrain it on unlabelled images without writing the BYOL loss, projection head and momentum update yourself. Do not adopt it if your downstream task is segmentation, where the README points at the separate pixel-level-contrastive-learning repository, or if you need the augmentation policy expressed as data rather than as an nn.Module. Before committing, verify that your chosen hidden_layer name matches an actual module in your model, and confirm that the batch sizes you can afford are large enough for the default SimCLR-derived augmentations to make sense.
Community notes