byol-pytorch

To implement the SimSiam variant, set use_momentum = False in the BYOL constructor. This removes the need for an exponential moving average target encoder. When using this mode, you do not need to call learner.update_moving_average() during the training loop.

import torch
from byol_pytorch import BYOL
from torchvision import models

resnet = models.resnet50(pretrained=True)

learner = BYOL(
    resnet,
    image_size = 256,
    hidden_layer = 'avgpool',
    use_momentum = False       # turn off momentum in the target encoder
)

opt = torch.optim.Adam(learner.parameters(), lr=3e-4)

def sample_unlabelled_images():
    return torch.randn(20, 3, 256, 256)

for _ in range(100):
    images = sample_unlabelled_images()
    loss = learner(images)
    opt.zero_grad()
    loss.backward()
    opt.step()

# save your improved network
torch.save(resnet.state_dict(), './improved-net.pt')

The BYOL class wraps an existing image-based neural network to enable self-supervised learning. You must specify the image_size and the hidden_layer (the name or index of the layer whose output serves as the latent representation).

During training, you must manually call learner.update_moving_average() after each optimizer step to update the target encoder's moving average.

import torch
from byol_pytorch import BYOL
from torchvision import models

resnet = models.resnet50(pretrained=True)

learner = BYOL(
    resnet,
    image_size = 256,
    hidden_layer = 'avgpool'
)

opt = torch.optim.Adam(learner.parameters(), lr=3e-4)

def sample_unlabelled_images():
    return torch.randn(20, 3, 256, 256)

for _ in range(100):
    images = sample_unlabelled_images()
    loss = learner(images)
    opt.zero_grad()
    loss.backward()
    opt.step()
    learner.update_moving_average() # update moving average of target encoder

# save your improved network
torch.save(resnet.state_dict(), './improved-net.pt')

For distributed training, use the BYOLTrainer class integrated with Hugging Face accelerate.

1. Configure your environment using accelerate config.
2. Initialize BYOLTrainer with your model, a Dataset, and training parameters like learning_rate, num_train_steps, batch_size, and checkpoint_every.
3. Launch your training script using accelerate launch <script.py>.

# 1. Setup configuration
$ accelerate config

# 2. Craft training script (e.g., ./train.py)
from torchvision import models
from byol_pytorch import (
    BYOL,
    BYOLTrainer,
    MockDataset
)

resnet = models.resnet50(pretrained = True)

dataset = MockDataset(256, 10000)

trainer = BYOLTrainer(
    resnet,
    dataset = dataset,
    image_size = 256,
    hidden_layer = 'avgpool',
    learning_rate = 3e-4,
    num_train_steps = 100_000,
    batch_size = 16,
    checkpoint_every = 1000     # improved model will be saved periodically to ./checkpoints folder
)

trainer()

# 3. Launch training
$ accelerate launch ./train.py

You can customize the BYOL training process using several keyword arguments:

• projection_size: The size of the projection.
• projection_hidden_size: The hidden dimension of the MLP for both the projection and prediction.
• moving_average_decay: The moving average decay factor for the target encoder.
• augment_fn: A custom augmentation function for the first view.
• augment_fn2: A custom augmentation function for the second view (useful for applying different probabilities of effects like Gaussian blur).

Example of custom hyperparameter and augmentation configuration:

# Custom hyperparameters
learner = BYOL(
    resnet,
    image_size = 256,
    hidden_layer = 'avgpool',
    projection_size = 256,
    projection_hidden_size = 4096,
    moving_average_decay = 0.99
)

# Custom augmentations
augment_fn = nn.Sequential(
    kornia.augmentation.RandomHorizontalFlip()
)

augment_fn2 = nn.Sequential(
    kornia.augmentation.RandomHorizontalFlip(),
    kornia.filters.GaussianBlur2d((3, 3), (1.5, 1.5))
)

learner = BYOL(
    resnet,
    image_size = 256,
    hidden_layer = -2,
    augment_fn = augment_fn,
    augment_fn2 = augment_fn2,
)

The BYOLTrainer constructor accepts several configuration options: