From 1ce32651f31d83784cb7b85d40caf886d4b8400b Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Wed, 6 Dec 2023 13:08:04 +0000 Subject: [PATCH 01/25] Adding SSL-EY --- examples/pytorch/ssley.py | 70 +++++++++ examples/pytorch_lightning/ssley.py | 67 +++++++++ .../pytorch_lightning_distributed/ssley.py | 77 ++++++++++ lightly/loss/ssley_loss.py | 140 ++++++++++++++++++ 4 files changed, 354 insertions(+) create mode 100644 examples/pytorch/ssley.py create mode 100644 examples/pytorch_lightning/ssley.py create mode 100644 examples/pytorch_lightning_distributed/ssley.py create mode 100644 lightly/loss/ssley_loss.py diff --git a/examples/pytorch/ssley.py b/examples/pytorch/ssley.py new file mode 100644 index 000000000..695fe92de --- /dev/null +++ b/examples/pytorch/ssley.py @@ -0,0 +1,70 @@ +import torch +import torchvision +from torch import nn + +## The projection head is the same as the Barlow Twins one +from lightly.loss import SSLEYLoss + +## The projection head is the same as the Barlow Twins one +from lightly.loss.ssley_loss import SSLEYLoss +from lightly.models.modules.heads import VICRegProjectionHead +from lightly.transforms.vicreg_transform import VICRegTransform + + +class SSLEY(nn.Module): + def __init__(self, backbone): + super().__init__() + self.backbone = backbone + self.projection_head = VICRegProjectionHead( + input_dim=512, + hidden_dim=2048, + output_dim=2048, + num_layers=2, + ) + + def forward(self, x): + x = self.backbone(x).flatten(start_dim=1) + z = self.projection_head(x) + return z + + +resnet = torchvision.models.resnet18() +backbone = nn.Sequential(*list(resnet.children())[:-1]) +model = SSLEY(backbone) + +device = "cuda" if torch.cuda.is_available() else "cpu" +model.to(device) + +transform = VICRegTransform(input_size=32) +dataset = torchvision.datasets.CIFAR10( + "datasets/cifar10", download=True, transform=transform +) +# or create a dataset from a folder containing images or videos: +# dataset = LightlyDataset("path/to/folder", transform=transform) + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=256, + shuffle=True, + drop_last=True, + num_workers=8, +) +criterion = SSLEYLoss() +optimizer = torch.optim.SGD(model.parameters(), lr=0.06) + +print("Starting Training") +for epoch in range(10): + total_loss = 0 + for batch in dataloader: + x0, x1 = batch[0] + x0 = x0.to(device) + x1 = x1.to(device) + z0 = model(x0) + z1 = model(x1) + loss = criterion(z0, z1) + total_loss += loss.detach() + loss.backward() + optimizer.step() + optimizer.zero_grad() + avg_loss = total_loss / len(dataloader) + print(f"epoch: {epoch:>02}, loss: {avg_loss:.5f}") diff --git a/examples/pytorch_lightning/ssley.py b/examples/pytorch_lightning/ssley.py new file mode 100644 index 000000000..6cc17de96 --- /dev/null +++ b/examples/pytorch_lightning/ssley.py @@ -0,0 +1,67 @@ +# Note: The model and training settings do not follow the reference settings +# from the paper. The settings are chosen such that the example can easily be +# run on a small dataset with a single GPU. + +import pytorch_lightning as pl +import torch +import torchvision +from torch import nn + +from lightly.loss.ssley_loss import SSLEYLoss + +## The projection head is the same as the Barlow Twins one +from lightly.models.modules.heads import VICRegProjectionHead +from lightly.transforms.vicreg_transform import VICRegTransform + + +class SSLEY(pl.LightningModule): + def __init__(self): + super().__init__() + resnet = torchvision.models.resnet18() + self.backbone = nn.Sequential(*list(resnet.children())[:-1]) + self.projection_head = VICRegProjectionHead( + input_dim=512, + hidden_dim=2048, + output_dim=2048, + num_layers=2, + ) + self.criterion = SSLEYLoss() + + def forward(self, x): + x = self.backbone(x).flatten(start_dim=1) + z = self.projection_head(x) + return z + + def training_step(self, batch, batch_index): + (x0, x1) = batch[0] + z0 = self.forward(x0) + z1 = self.forward(x1) + loss = self.criterion(z0, z1) + return loss + + def configure_optimizers(self): + optim = torch.optim.SGD(self.parameters(), lr=0.06) + return optim + + +model = SSLEY() + +transform = VICRegTransform(input_size=32) +dataset = torchvision.datasets.CIFAR10( + "datasets/cifar10", download=True, transform=transform +) +# or create a dataset from a folder containing images or videos: +# dataset = LightlyDataset("path/to/folder", transform=transform) + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=256, + shuffle=True, + drop_last=True, + num_workers=8, +) + +accelerator = "gpu" if torch.cuda.is_available() else "cpu" + +trainer = pl.Trainer(max_epochs=10, devices=1, accelerator=accelerator) +trainer.fit(model=model, train_dataloaders=dataloader) diff --git a/examples/pytorch_lightning_distributed/ssley.py b/examples/pytorch_lightning_distributed/ssley.py new file mode 100644 index 000000000..3f034adcf --- /dev/null +++ b/examples/pytorch_lightning_distributed/ssley.py @@ -0,0 +1,77 @@ +# Note: The model and training settings do not follow the reference settings +# from the paper. The settings are chosen such that the example can easily be +# run on a small dataset with a single GPU. + +import pytorch_lightning as pl +import torch +import torchvision +from torch import nn + +from lightly.loss import SSLEYLoss + +## The projection head is the same as the Barlow Twins one +from lightly.models.modules.heads import VICRegProjectionHead +from lightly.transforms.vicreg_transform import VICRegTransform + + +class SSLEY(pl.LightningModule): + def __init__(self): + super().__init__() + resnet = torchvision.models.resnet18() + self.backbone = nn.Sequential(*list(resnet.children())[:-1]) + self.projection_head = VICRegProjectionHead( + input_dim=512, + hidden_dim=2048, + output_dim=2048, + num_layers=2, + ) + + # enable gather_distributed to gather features from all gpus + # before calculating the loss + self.criterion = SSLEYLoss(gather_distributed=True) + + def forward(self, x): + x = self.backbone(x).flatten(start_dim=1) + z = self.projection_head(x) + return z + + def training_step(self, batch, batch_index): + (x0, x1) = batch[0] + z0 = self.forward(x0) + z1 = self.forward(x1) + loss = self.criterion(z0, z1) + return loss + + def configure_optimizers(self): + optim = torch.optim.SGD(self.parameters(), lr=0.06) + return optim + + +model = SSLEY() + +transform = VICRegTransform(input_size=32) +dataset = torchvision.datasets.CIFAR10( + "datasets/cifar10", download=True, transform=transform +) +# or create a dataset from a folder containing images or videos: +# dataset = LightlyDataset("path/to/folder", transform=transform) + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=256, + shuffle=True, + drop_last=True, + num_workers=8, +) + +# Train with DDP and use Synchronized Batch Norm for a more accurate batch norm +# calculation. Distributed sampling is also enabled with replace_sampler_ddp=True. +trainer = pl.Trainer( + max_epochs=10, + devices="auto", + accelerator="gpu", + strategy="ddp", + sync_batchnorm=True, + use_distributed_sampler=True, # or replace_sampler_ddp=True for PyTorch Lightning <2.0 +) +trainer.fit(model=model, train_dataloaders=dataloader) diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py new file mode 100644 index 000000000..74dc1e1ee --- /dev/null +++ b/lightly/loss/ssley_loss.py @@ -0,0 +1,140 @@ +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import Tensor + +from lightly.utils.dist import gather + + +class SSLEYLoss(torch.nn.Module): + """Implementation of the SSL-EY loss [0]. + + - [0] Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2022, https://arxiv.org/abs/2310.01012 + + Attributes: + gather_distributed: + If True then the cross-correlation matrices from all gpus are gathered and + summed before the loss calculation. + eps: + Epsilon for numerical stability. + + Examples: + + >>> # initialize loss function + >>> loss_fn = SSLEYLoss() + >>> + >>> # generate two random transforms of images + >>> t0 = transforms(images) + >>> t1 = transforms(images) + >>> + >>> # feed through model + >>> out0, out1 = model(t0, t1) + >>> + >>> # calculate loss + >>> loss = loss_fn(out0, out1) + """ + + def __init__( + self, + gather_distributed: bool = False, + eps=0.0001, + ): + super(SSLEYLoss, self).__init__() + if gather_distributed and not dist.is_available(): + raise ValueError( + "gather_distributed is True but torch.distributed is not available. " + "Please set gather_distributed=False or install a torch version with " + "distributed support." + ) + + self.lambda_param = lambda_param + self.mu_param = mu_param + self.nu_param = nu_param + self.gather_distributed = gather_distributed + self.eps = eps + + def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: + """Returns SSL-EY loss. + + Args: + z_a: + Tensor with shape (batch_size, ..., dim). + z_b: + Tensor with shape (batch_size, ..., dim). + """ + assert ( + z_a.shape[0] > 1 and z_b.shape[0] > 1 + ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}" + assert ( + z_a.shape == z_b.shape + ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}." + + # invariance term of the loss + inv_loss = invariance_loss(x=z_a, y=z_b) + + # gather all batches + if self.gather_distributed and dist.is_initialized(): + world_size = dist.get_world_size() + if world_size > 1: + z_a = torch.cat(gather(z_a), dim=0) + z_b = torch.cat(gather(z_b), dim=0) + + var_loss = 0.5 * ( + variance_loss(x=z_a, eps=self.eps) + variance_loss(x=z_b, eps=self.eps) + ) + cov_loss = covariance_loss(x=z_a) + covariance_loss(x=z_b) + + loss = ( + self.lambda_param * inv_loss + + self.mu_param * var_loss + + self.nu_param * cov_loss + ) + return loss + + +def invariance_loss(x: Tensor, y: Tensor) -> Tensor: + """Returns SSL-EY invariance loss. + + Args: + x: + Tensor with shape (batch_size, ..., dim). + y: + Tensor with shape (batch_size, ..., dim). + """ + return F.mse_loss(x, y) + + +def variance_loss(x: Tensor, eps: float = 0.0001) -> Tensor: + """Returns SSL-EY variance loss. + + Args: + x: + Tensor with shape (batch_size, ..., dim). + eps: + Epsilon for numerical stability. + """ + x = x - x.mean(dim=0) + std = torch.sqrt(x.var(dim=0) + eps) + loss = torch.mean(F.relu(1.0 - std)) + return loss + + +def covariance_loss(x: Tensor) -> Tensor: + """Returns SSL-EY covariance loss. + + Generalized version of the covariance loss with support for tensors with more than + two dimensions. + + Args: + x: + Tensor with shape (batch_size, ..., dim). + """ + x = x - x.mean(dim=0) + batch_size = x.size(0) + dim = x.size(-1) + # nondiag_mask has shape (dim, dim) with 1s on all non-diagonal entries. + nondiag_mask = ~torch.eye(dim, device=x.device, dtype=torch.bool) + # cov has shape (..., dim, dim) + cov = torch.einsum("b...c,b...d->...cd", x, x) / (batch_size - 1) + loss = cov[..., nondiag_mask].pow(2).sum(-1) / dim + return loss.mean() From c8943b16ec71af2e508f01828dbc1342021d5a72 Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Wed, 6 Dec 2023 13:14:40 +0000 Subject: [PATCH 02/25] Adding SSL-EY --- docs/source/examples/ssley.rst | 48 ++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 docs/source/examples/ssley.rst diff --git a/docs/source/examples/ssley.rst b/docs/source/examples/ssley.rst new file mode 100644 index 000000000..8a449f692 --- /dev/null +++ b/docs/source/examples/ssley.rst @@ -0,0 +1,48 @@ +.. _ssley: + +SSL-EY +======= + +SSL-EY is a method that explicitly +avoids the collapse problem with a simple regularization term on the variance of the embeddings along each dimension individually. It inherits the model structure from +`Barlow Twins, 2022 `_ changing the loss. Doing so allows the stabilization of the training and leads to performance improvements. + +Reference: + `Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2023 `_ + + +.. tabs:: + .. tab:: PyTorch + + This example can be run from the command line with:: + + python lightly/examples/pytorch/ssley.py + + .. literalinclude:: ../../../examples/pytorch/ssley.py + + .. tab:: Lightning + + This example can be run from the command line with:: + + python lightly/examples/pytorch_lightning/ssley.py + + .. literalinclude:: ../../../examples/pytorch_lightning/ssley.py + + .. tab:: Lightning Distributed + + This example runs on multiple gpus using Distributed Data Parallel (DDP) + training with Pytorch Lightning. At least one GPU must be available on + the system. The example can be run from the command line with:: + + python lightly/examples/pytorch_lightning_distributed/ssley.py + + The model differs in the following ways from the non-distributed + implementation: + + - Distributed Data Parallel is enabled + - Distributed Sampling is used in the dataloader + + Distributed Sampling makes sure that each distributed process sees only + a subset of the data. + + .. literalinclude:: ../../../examples/pytorch_lightning_distributed/ssley.py \ No newline at end of file From 410d9714f6cdcf1efca308786f80a219f3d509cc Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Wed, 6 Dec 2023 13:24:09 +0000 Subject: [PATCH 03/25] Adding SSL-EY to docs --- README.md | 22 +-- benchmarks/imagenet/resnet50/ssley.py | 127 ++++++++++++++++++ .../benchmarks/imagenette_benchmark.py | 41 ++++++ 3 files changed, 180 insertions(+), 10 deletions(-) create mode 100644 benchmarks/imagenet/resnet50/ssley.py diff --git a/README.md b/README.md index 70afc9418..77e10af35 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,7 @@ and PyTorch Lightning distributed examples for all models to kickstart your proj - SimMIM, 2021 [paper](https://arxiv.org/abs/2111.09886) [docs](https://docs.lightly.ai/self-supervised-learning/examples/simmim.html) - SimSiam, 2021 [paper](https://arxiv.org/abs/2011.10566) [docs](https://docs.lightly.ai/self-supervised-learning/examples/simsiam.html) - SMoG, 2022 [paper](https://arxiv.org/abs/2207.06167) [docs](https://docs.lightly.ai/self-supervised-learning/examples/smog.html) +- SSL-EY, 2023 [paper](https://arxiv.org/abs/2310.01012) [docs](https://docs.lightly.ai/self-supervised-learning/examples/ssley.html) - SwaV, 2020 [paper](https://arxiv.org/abs/2006.09882) [docs](https://docs.lightly.ai/self-supervised-learning/examples/swav.html) - TiCo, 2022 [paper](https://arxiv.org/abs/2206.10698) [docs](https://docs.lightly.ai/self-supervised-learning/examples/tico.html) - VICReg, 2022 [paper](https://arxiv.org/abs/2105.04906) [docs](https://docs.lightly.ai/self-supervised-learning/examples/vicreg.html) @@ -278,16 +279,17 @@ tuned for maximum accuracy. For detailed results and more information about the See the [benchmarking scripts](./benchmarks/imagenet/resnet50/) for details. -| Model | Backbone | Batch Size | Epochs | Linear Top1 | Finetune Top1 | kNN Top1 | Tensorboard | Checkpoint | -|----------------|----------|------------|--------|-------------|---------------|----------|-------------|------------| -| BarlowTwins | Res50 | 256 | 100 | 62.9 | 72.6 | 45.6 | [link](https://tensorboard.dev/experiment/NxyNRiQsQjWZ82I9b0PvKg/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_barlowtwins_2023-08-18_00-11-03/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | -| BYOL | Res50 | 256 | 100 | 62.4 | 74.0 | 45.6 | [link](https://tensorboard.dev/experiment/Z0iG2JLaTJe5nuBD7DK1bg) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_byol_2023-07-10_10-37-32/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | -| DINO | Res50 | 128 | 100 | 68.2 | 72.5 | 49.9 | [link](https://tensorboard.dev/experiment/DvKHX9sNSWWqDrRksllPLA) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dino_2023-06-06_13-59-48/pretrain/version_0/checkpoints/epoch%3D99-step%3D1000900.ckpt) | -| SimCLR* | Res50 | 256 | 100 | 63.2 | 73.9 | 44.8 | [link](https://tensorboard.dev/experiment/Ugol97adQdezgcVibDYMMA) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_simclr_2023-06-22_09-11-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | -| SimCLR* + DCL | Res50 | 256 | 100 | 65.1 | 73.5 | 49.6 | [link](https://tensorboard.dev/experiment/k4ZonZ77QzmBkc0lXswQlg/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dcl_2023-07-04_16-51-40/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | -| SimCLR* + DCLW | Res50 | 256 | 100 | 64.5 | 73.2 | 48.5 | [link](https://tensorboard.dev/experiment/TrALnpwFQ4OkZV3uvaX7wQ/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dclw_2023-07-07_14-57-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | -| SwAV | Res50 | 256 | 100 | 67.2 | 75.4 | 49.5 | [link](https://tensorboard.dev/experiment/Ipx4Oxl5Qkqm5Sl5kWyKKg) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_swav_2023-05-25_08-29-14/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | -| VICReg | Res50 | 256 | 100 | 63.0 | 73.7 | 46.3 | [link](https://tensorboard.dev/experiment/qH5uywJbTJSzgCEfxc7yUw) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_vicreg_2023-09-11_10-53-08/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | +| Model | Backbone | Batch Size | Epochs | Linear Top1 | Finetune Top1 | kNN Top1 | Tensorboard | Checkpoint | +|----------------|----------|------------|--------|-------------|---------------|----------|----------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| BarlowTwins | Res50 | 256 | 100 | 62.9 | 72.6 | 45.6 | [link](https://tensorboard.dev/experiment/NxyNRiQsQjWZ82I9b0PvKg/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_barlowtwins_2023-08-18_00-11-03/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | +| BYOL | Res50 | 256 | 100 | 62.4 | 74.0 | 45.6 | [link](https://tensorboard.dev/experiment/Z0iG2JLaTJe5nuBD7DK1bg) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_byol_2023-07-10_10-37-32/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | +| DINO | Res50 | 128 | 100 | 68.2 | 72.5 | 49.9 | [link](https://tensorboard.dev/experiment/DvKHX9sNSWWqDrRksllPLA) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dino_2023-06-06_13-59-48/pretrain/version_0/checkpoints/epoch%3D99-step%3D1000900.ckpt) | +| SimCLR* | Res50 | 256 | 100 | 63.2 | 73.9 | 44.8 | [link](https://tensorboard.dev/experiment/Ugol97adQdezgcVibDYMMA) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_simclr_2023-06-22_09-11-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | +| SimCLR* + DCL | Res50 | 256 | 100 | 65.1 | 73.5 | 49.6 | [link](https://tensorboard.dev/experiment/k4ZonZ77QzmBkc0lXswQlg/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dcl_2023-07-04_16-51-40/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | +| SimCLR* + DCLW | Res50 | 256 | 100 | 64.5 | 73.2 | 48.5 | [link](https://tensorboard.dev/experiment/TrALnpwFQ4OkZV3uvaX7wQ/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dclw_2023-07-07_14-57-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | +| SSL-EY | Res50 | 256 | 100 | TODO | TODO | TODO | [link](TODO)| [link](TODO) | +| SwAV | Res50 | 256 | 100 | 67.2 | 75.4 | 49.5 | [link](https://tensorboard.dev/experiment/Ipx4Oxl5Qkqm5Sl5kWyKKg) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_swav_2023-05-25_08-29-14/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | +| VICReg | Res50 | 256 | 100 | 63.0 | 73.7 | 46.3 | [link](https://tensorboard.dev/experiment/qH5uywJbTJSzgCEfxc7yUw) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_vicreg_2023-09-11_10-53-08/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) | *\*We use square root learning rate scaling instead of linear scaling as it yields better results for smaller batch sizes. See Appendix B.1 in the [SimCLR paper](https://arxiv.org/abs/2002.05709).* diff --git a/benchmarks/imagenet/resnet50/ssley.py b/benchmarks/imagenet/resnet50/ssley.py new file mode 100644 index 000000000..1c45a75ee --- /dev/null +++ b/benchmarks/imagenet/resnet50/ssley.py @@ -0,0 +1,127 @@ +from typing import List, Tuple + +import torch +from pytorch_lightning import LightningModule +from torch import Tensor +from torch.nn import Identity +from torchvision.models import resnet50 + +from lightly.loss.ssley_loss import SSLEYLoss +from lightly.models.modules.heads import VICRegProjectionHead +from lightly.models.utils import get_weight_decay_parameters +from lightly.transforms.ssley_transform import VICRegTransform +from lightly.utils.benchmarking import OnlineLinearClassifier +from lightly.utils.lars import LARS +from lightly.utils.scheduler import CosineWarmupScheduler + + +class SSLEY(LightningModule): + def __init__(self, batch_size_per_device: int, num_classes: int) -> None: + super().__init__() + self.save_hyperparameters() + self.batch_size_per_device = batch_size_per_device + + resnet = resnet50() + resnet.fc = Identity() # Ignore classification head + self.backbone = resnet + self.projection_head = VICRegProjectionHead(num_layers=2) + self.criterion = SSLEYLoss() + + self.online_classifier = OnlineLinearClassifier(num_classes=num_classes) + + def forward(self, x: Tensor) -> Tensor: + return self.backbone(x) + + def training_step( + self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int + ) -> Tensor: + views, targets = batch[0], batch[1] + features = self.forward(torch.cat(views)).flatten(start_dim=1) + z = self.projection_head(features) + z_a, z_b = z.chunk(len(views)) + loss = self.criterion(z_a=z_a, z_b=z_b) + self.log( + "train_loss", loss, prog_bar=True, sync_dist=True, batch_size=len(targets) + ) + + # Online linear evaluation. + cls_loss, cls_log = self.online_classifier.training_step( + (features.detach(), targets.repeat(len(views))), batch_idx + ) + + self.log_dict(cls_log, sync_dist=True, batch_size=len(targets)) + return loss + cls_loss + + def validation_step( + self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int + ) -> Tensor: + images, targets = batch[0], batch[1] + features = self.forward(images).flatten(start_dim=1) + cls_loss, cls_log = self.online_classifier.validation_step( + (features.detach(), targets), batch_idx + ) + self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets)) + return cls_loss + + def configure_optimizers(self): + # Don't use weight decay for batch norm, bias parameters, and classification + # head to improve performance. + params, params_no_weight_decay = get_weight_decay_parameters( + [self.backbone, self.projection_head] + ) + global_batch_size = self.batch_size_per_device * self.trainer.world_size + base_lr = _get_base_learning_rate(global_batch_size=global_batch_size) + optimizer = LARS( + [ + {"name": "ssley", "params": params}, + { + "name": "ssley_no_weight_decay", + "params": params_no_weight_decay, + "weight_decay": 0.0, + }, + { + "name": "online_classifier", + "params": self.online_classifier.parameters(), + "weight_decay": 0.0, + }, + ], + # Linear learning rate scaling with a base learning rate of 0.2. + # See https://arxiv.org/pdf/2105.04906.pdf for details. + lr=base_lr * global_batch_size / 256, + momentum=0.9, + weight_decay=1e-6, + ) + scheduler = { + "scheduler": CosineWarmupScheduler( + optimizer=optimizer, + warmup_epochs=( + self.trainer.estimated_stepping_batches + / self.trainer.max_epochs + * 10 + ), + max_epochs=self.trainer.estimated_stepping_batches, + end_value=0.01, # Scale base learning rate from 0.2 to 0.002. + ), + "interval": "step", + } + return [optimizer], [scheduler] + + +# SSLEY transform +transform = VICRegTransform() + + +def _get_base_learning_rate(global_batch_size: int) -> float: + """Returns the base learning rate for training 100 epochs with a given batch size. + + This follows section C.4 in https://arxiv.org/pdf/2105.04906.pdf. + + """ + if global_batch_size == 128: + return 0.8 + elif global_batch_size == 256: + return 0.5 + elif global_batch_size == 512: + return 0.4 + else: + return 0.3 diff --git a/docs/source/getting_started/benchmarks/imagenette_benchmark.py b/docs/source/getting_started/benchmarks/imagenette_benchmark.py index b5c5e8e71..af2759ed0 100644 --- a/docs/source/getting_started/benchmarks/imagenette_benchmark.py +++ b/docs/source/getting_started/benchmarks/imagenette_benchmark.py @@ -29,6 +29,7 @@ | SimCLR | 256 | 200 | 0.835 | 49.7 Min | 3.7 GByte | | SimMIM (ViT-B32) | 256 | 200 | 0.315 | 115.5 Min | 9.7 GByte | | SimSiam | 256 | 200 | 0.752 | 58.2 Min | 3.9 GByte | +| SSL-EY | 256 | 200 | TO-DO | TO-DO | TO-DO GByte | | SwaV | 256 | 200 | 0.861 | 73.3 Min | 6.4 GByte | | SwaVQueue | 256 | 200 | 0.827 | 72.6 Min | 6.4 GByte | | SMoG | 256 | 200 | 0.663 | 58.7 Min | 2.6 GByte | @@ -50,6 +51,7 @@ | SimCLR | 256 | 800 | 0.889 | 193.5 Min | 3.7 GByte | | SimMIM (ViT-B32) | 256 | 800 | 0.343 | 446.5 Min | 9.7 GByte | | SimSiam | 256 | 800 | 0.872 | 206.4 Min | 3.9 GByte | +| SSL-EY | 256 | 800 | TO-DO | TO-DO | TO-DO GByte | | SwaV | 256 | 800 | 0.902 | 283.2 Min | 6.4 GByte | | SwaVQueue | 256 | 800 | 0.890 | 282.7 Min | 6.4 GByte | | SMoG | 256 | 800 | 0.788 | 232.1 Min | 2.6 GByte | @@ -81,6 +83,7 @@ NegativeCosineSimilarity, NTXentLoss, PMSNLoss, + SSLEYLoss, SwaVLoss, TiCoLoss, VICRegLLoss, @@ -267,6 +270,7 @@ def create_dataset_train_ssl(model): SimCLRModel: simclr_transform, SimMIMModel: simmim_transform, SimSiamModel: simsiam_transform, + SSL_EYModel: vicreg_transform, SwaVModel: swav_transform, SwaVQueueModel: swav_transform, SMoGModel: smog_transform, @@ -1165,6 +1169,42 @@ def configure_optimizers(self): return [optim], [cosine_scheduler] +class SSLEYModel(BenchmarkModule): + def __init__(self, dataloader_kNN, num_classes): + super().__init__(dataloader_kNN, num_classes) + # create a ResNet backbone and remove the classification head + resnet = torchvision.models.resnet18() + self.backbone = nn.Sequential(*list(resnet.children())[:-1]) + self.projection_head = heads.BarlowTwinsProjectionHead(512, 2048, 2048) + self.criterion = SSLEYLoss() + self.warmup_epochs = 40 if max_epochs >= 800 else 20 + + def forward(self, x): + x = self.backbone(x).flatten(start_dim=1) + z = self.projection_head(x) + return z + + def training_step(self, batch, batch_index): + (x0, x1), _, _ = batch + z0 = self.forward(x0) + z1 = self.forward(x1) + loss = self.criterion(z0, z1) + return loss + + def configure_optimizers(self): + # Training diverges without LARS + optim = LARS( + self.parameters(), + lr=0.3 * lr_factor, + weight_decay=1e-4, + momentum=0.9, + ) + cosine_scheduler = scheduler.CosineWarmupScheduler( + optim, self.warmup_epochs, max_epochs + ) + return [optim], [cosine_scheduler] + + class VICRegModel(BenchmarkModule): def __init__(self, dataloader_kNN, num_classes): super().__init__(dataloader_kNN, num_classes) @@ -1410,6 +1450,7 @@ def configure_optimizers(self): SimCLRModel, # SimMIMModel, # disabled by default because SimMIM uses larger images with size 224 SimSiamModel, + SSLEYModel, SwaVModel, SwaVQueueModel, SMoGModel, From fbd64389f7fdbfc6cb7e2c7f5604fb18edf4530e Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Wed, 6 Dec 2023 13:39:07 +0000 Subject: [PATCH 04/25] Adding SSL-EY to docs --- benchmarks/imagenet/resnet50/main.py | 2 + docs/source/lightly.loss.rst | 3 ++ tests/loss/test_SSLEYLoss.py | 67 ++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+) create mode 100644 tests/loss/test_SSLEYLoss.py diff --git a/benchmarks/imagenet/resnet50/main.py b/benchmarks/imagenet/resnet50/main.py index 8eb3e7f65..219060719 100644 --- a/benchmarks/imagenet/resnet50/main.py +++ b/benchmarks/imagenet/resnet50/main.py @@ -12,6 +12,7 @@ import knn_eval import linear_eval import simclr +import ssley import swav import torch import vicreg @@ -58,6 +59,7 @@ "dclw": {"model": dclw.DCLW, "transform": dclw.transform}, "dino": {"model": dino.DINO, "transform": dino.transform}, "simclr": {"model": simclr.SimCLR, "transform": simclr.transform}, + "ssley": {"model": ssley.SSLEY, "transform": ssley.transform}, "swav": {"model": swav.SwAV, "transform": swav.transform}, "vicreg": {"model": vicreg.VICReg, "transform": vicreg.transform}, } diff --git a/docs/source/lightly.loss.rst b/docs/source/lightly.loss.rst index 38e304b64..2976cca51 100644 --- a/docs/source/lightly.loss.rst +++ b/docs/source/lightly.loss.rst @@ -41,6 +41,9 @@ lightly.loss .. autoclass:: lightly.loss.regularizer.co2.CO2Regularizer :members: +.. autoclass:: lightly.loss.ssley_loss.SSLEYLoss + :members: + .. autoclass:: lightly.loss.swav_loss.SwaVLoss :members: diff --git a/tests/loss/test_SSLEYLoss.py b/tests/loss/test_SSLEYLoss.py new file mode 100644 index 000000000..a803e5885 --- /dev/null +++ b/tests/loss/test_SSLEYLoss.py @@ -0,0 +1,67 @@ +import unittest + +import pytest +import torch +import torch.nn.functional as F +from pytest_mock import MockerFixture +from torch import Tensor +from torch import distributed as dist + +from lightly.loss import SSLEYLoss + + +class TestSSLEYLoss: + def test__gather_distributed(self, mocker: MockerFixture) -> None: + mock_is_available = mocker.patch.object(dist, "is_available", return_value=True) + SSLEYLoss(gather_distributed=True) + mock_is_available.assert_called_once() + + def test__gather_distributed_dist_not_available( + self, mocker: MockerFixture + ) -> None: + mock_is_available = mocker.patch.object( + dist, "is_available", return_value=False + ) + with pytest.raises(ValueError): + SSLEYLoss(gather_distributed=True) + mock_is_available.assert_called_once() + + +class TestSSLEYLossUnitTest(unittest.TestCase): + # Old tests in unittest style, please add new tests to TestSSLEYLoss using pytest. + def test_forward_pass(self): + loss = SSLEYLoss() + for bsz in range(2, 4): + x0 = torch.randn((bsz, 32)) + x1 = torch.randn((bsz, 32)) + + # symmetry + l1 = loss(x0, x1) + l2 = loss(x1, x0) + self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0) + + @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available") + def test_forward_pass_cuda(self): + loss = SSLEYLoss() + for bsz in range(2, 4): + x0 = torch.randn((bsz, 32)).cuda() + x1 = torch.randn((bsz, 32)).cuda() + + # symmetry + l1 = loss(x0, x1) + l2 = loss(x1, x0) + self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0) + + def test_forward_pass__error_batch_size_1(self): + loss = SSLEYLoss() + x0 = torch.randn((1, 32)) + x1 = torch.randn((1, 32)) + with self.assertRaises(AssertionError): + loss(x0, x1) + + def test_forward_pass__error_different_shapes(self): + loss = SSLEYLoss() + x0 = torch.randn((2, 32)) + x1 = torch.randn((2, 16)) + with self.assertRaises(AssertionError): + loss(x0, x1) From d41c0f718d28dc78db2386e4ce9df5879cd7e458 Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Wed, 6 Dec 2023 14:14:05 +0000 Subject: [PATCH 05/25] Adding SSL-EY to tests --- lightly/loss/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lightly/loss/__init__.py b/lightly/loss/__init__.py index 613707abe..51ff573c3 100644 --- a/lightly/loss/__init__.py +++ b/lightly/loss/__init__.py @@ -9,6 +9,7 @@ from lightly.loss.negative_cosine_similarity import NegativeCosineSimilarity from lightly.loss.ntx_ent_loss import NTXentLoss from lightly.loss.pmsn_loss import PMSNCustomLoss, PMSNLoss +from lightly.loss.ssley_loss import SSLEYLoss from lightly.loss.swav_loss import SwaVLoss from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss from lightly.loss.tico_loss import TiCoLoss From 6dbb3014eb1c044a2c3b5f37d12d333b1b0aaac2 Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Fri, 8 Dec 2023 11:30:21 +0000 Subject: [PATCH 06/25] Adding SSL-EY to tests --- lightly/loss/ssley_loss.py | 74 ++++---------------------------------- 1 file changed, 6 insertions(+), 68 deletions(-) diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py index 74dc1e1ee..e10c37560 100644 --- a/lightly/loss/ssley_loss.py +++ b/lightly/loss/ssley_loss.py @@ -47,9 +47,6 @@ def __init__( "distributed support." ) - self.lambda_param = lambda_param - self.mu_param = mu_param - self.nu_param = nu_param self.gather_distributed = gather_distributed self.eps = eps @@ -62,15 +59,6 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: z_b: Tensor with shape (batch_size, ..., dim). """ - assert ( - z_a.shape[0] > 1 and z_b.shape[0] > 1 - ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}" - assert ( - z_a.shape == z_b.shape - ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}." - - # invariance term of the loss - inv_loss = invariance_loss(x=z_a, y=z_b) # gather all batches if self.gather_distributed and dist.is_initialized(): @@ -79,62 +67,12 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: z_a = torch.cat(gather(z_a), dim=0) z_b = torch.cat(gather(z_b), dim=0) - var_loss = 0.5 * ( - variance_loss(x=z_a, eps=self.eps) + variance_loss(x=z_b, eps=self.eps) - ) - cov_loss = covariance_loss(x=z_a) + covariance_loss(x=z_b) - - loss = ( - self.lambda_param * inv_loss - + self.mu_param * var_loss - + self.nu_param * cov_loss - ) - return loss - - -def invariance_loss(x: Tensor, y: Tensor) -> Tensor: - """Returns SSL-EY invariance loss. - - Args: - x: - Tensor with shape (batch_size, ..., dim). - y: - Tensor with shape (batch_size, ..., dim). - """ - return F.mse_loss(x, y) - + z_a = z_a - z_a.mean(dim=0) + z_b = z_b - z_b.mean(dim=0) -def variance_loss(x: Tensor, eps: float = 0.0001) -> Tensor: - """Returns SSL-EY variance loss. - - Args: - x: - Tensor with shape (batch_size, ..., dim). - eps: - Epsilon for numerical stability. - """ - x = x - x.mean(dim=0) - std = torch.sqrt(x.var(dim=0) + eps) - loss = torch.mean(F.relu(1.0 - std)) - return loss + C = 2*(z_a.T @ z_b) / (self.args.batch_size - 1) + V = (z_a.T @ z_a) / (self.args.batch_size - 1) + (z_b.T @ z_b) / (self.args.batch_size - 1) + loss = torch.trace(C)-torch.trace(V@V) -def covariance_loss(x: Tensor) -> Tensor: - """Returns SSL-EY covariance loss. - - Generalized version of the covariance loss with support for tensors with more than - two dimensions. - - Args: - x: - Tensor with shape (batch_size, ..., dim). - """ - x = x - x.mean(dim=0) - batch_size = x.size(0) - dim = x.size(-1) - # nondiag_mask has shape (dim, dim) with 1s on all non-diagonal entries. - nondiag_mask = ~torch.eye(dim, device=x.device, dtype=torch.bool) - # cov has shape (..., dim, dim) - cov = torch.einsum("b...c,b...d->...cd", x, x) / (batch_size - 1) - loss = cov[..., nondiag_mask].pow(2).sum(-1) / dim - return loss.mean() + return loss From ea0a3a079585e344875a0355c3207535406aa778 Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Fri, 8 Dec 2023 11:41:32 +0000 Subject: [PATCH 07/25] Adding SSL-EY to tests --- lightly/loss/ssley_loss.py | 15 ++++++++++----- tests/loss/test_SSLEYLoss.py | 2 -- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py index e10c37560..d7f015d6b 100644 --- a/lightly/loss/ssley_loss.py +++ b/lightly/loss/ssley_loss.py @@ -59,7 +59,12 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: z_b: Tensor with shape (batch_size, ..., dim). """ - + assert ( + z_a.shape[0] > 1 and z_b.shape[0] > 1 + ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}" + assert ( + z_a.shape == z_b.shape + ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}." # gather all batches if self.gather_distributed and dist.is_initialized(): world_size = dist.get_world_size() @@ -69,10 +74,10 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: z_a = z_a - z_a.mean(dim=0) z_b = z_b - z_b.mean(dim=0) + batch_size = z_a.size(0) + C = 2*(z_a.T @ z_b) / (batch_size - 1) + V = (z_a.T @ z_a) / (batch_size - 1) + (z_b.T @ z_b) / (batch_size - 1) - C = 2*(z_a.T @ z_b) / (self.args.batch_size - 1) - V = (z_a.T @ z_a) / (self.args.batch_size - 1) + (z_b.T @ z_b) / (self.args.batch_size - 1) - - loss = torch.trace(C)-torch.trace(V@V) + loss = -2 * torch.trace(C) + torch.trace(V @ V) return loss diff --git a/tests/loss/test_SSLEYLoss.py b/tests/loss/test_SSLEYLoss.py index a803e5885..0ee7b1e2b 100644 --- a/tests/loss/test_SSLEYLoss.py +++ b/tests/loss/test_SSLEYLoss.py @@ -2,9 +2,7 @@ import pytest import torch -import torch.nn.functional as F from pytest_mock import MockerFixture -from torch import Tensor from torch import distributed as dist from lightly.loss import SSLEYLoss From 3aae9fb5a6b2514f5ec52f817113488bd41e6f8a Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Fri, 15 Dec 2023 08:19:55 -0600 Subject: [PATCH 08/25] Adding SSL-EY to tests --- benchmarks/imagenet/resnet50/ssley.py | 2 +- lightly/loss/ssley_loss.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/benchmarks/imagenet/resnet50/ssley.py b/benchmarks/imagenet/resnet50/ssley.py index 1c45a75ee..b789b0c5d 100644 --- a/benchmarks/imagenet/resnet50/ssley.py +++ b/benchmarks/imagenet/resnet50/ssley.py @@ -9,7 +9,7 @@ from lightly.loss.ssley_loss import SSLEYLoss from lightly.models.modules.heads import VICRegProjectionHead from lightly.models.utils import get_weight_decay_parameters -from lightly.transforms.ssley_transform import VICRegTransform +from lightly.transforms.vicreg_transform import VICRegTransform from lightly.utils.benchmarking import OnlineLinearClassifier from lightly.utils.lars import LARS from lightly.utils.scheduler import CosineWarmupScheduler diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py index d7f015d6b..211621791 100644 --- a/lightly/loss/ssley_loss.py +++ b/lightly/loss/ssley_loss.py @@ -1,7 +1,5 @@ import torch import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor from lightly.utils.dist import gather From 7ebd3952351c54ef0d8b0e7d7ea8fe4b91ec23b2 Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Fri, 15 Dec 2023 08:24:48 -0600 Subject: [PATCH 09/25] Adding SSL-EY to tests --- examples/pytorch/ssley.py | 5 ----- lightly/loss/ssley_loss.py | 6 +++--- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/examples/pytorch/ssley.py b/examples/pytorch/ssley.py index 695fe92de..878737b56 100644 --- a/examples/pytorch/ssley.py +++ b/examples/pytorch/ssley.py @@ -1,11 +1,6 @@ import torch import torchvision from torch import nn - -## The projection head is the same as the Barlow Twins one -from lightly.loss import SSLEYLoss - -## The projection head is the same as the Barlow Twins one from lightly.loss.ssley_loss import SSLEYLoss from lightly.models.modules.heads import VICRegProjectionHead from lightly.transforms.vicreg_transform import VICRegTransform diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py index 211621791..b6023bd13 100644 --- a/lightly/loss/ssley_loss.py +++ b/lightly/loss/ssley_loss.py @@ -58,10 +58,10 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: Tensor with shape (batch_size, ..., dim). """ assert ( - z_a.shape[0] > 1 and z_b.shape[0] > 1 + z_a.shape[0] > 1 and z_b.shape[0] > 1 ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}" assert ( - z_a.shape == z_b.shape + z_a.shape == z_b.shape ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}." # gather all batches if self.gather_distributed and dist.is_initialized(): @@ -73,7 +73,7 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: z_a = z_a - z_a.mean(dim=0) z_b = z_b - z_b.mean(dim=0) batch_size = z_a.size(0) - C = 2*(z_a.T @ z_b) / (batch_size - 1) + C = 2 * (z_a.T @ z_b) / (batch_size - 1) V = (z_a.T @ z_a) / (batch_size - 1) + (z_b.T @ z_b) / (batch_size - 1) loss = -2 * torch.trace(C) + torch.trace(V @ V) From 0a08b0be73f2fbabda3c7be2b79b6887e7c52409 Mon Sep 17 00:00:00 2001 From: jameschapman19 Date: Fri, 15 Dec 2023 08:29:40 -0600 Subject: [PATCH 10/25] Adding SSL-EY to tests --- lightly/transforms/ssley_transform.py | 117 ++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 lightly/transforms/ssley_transform.py diff --git a/lightly/transforms/ssley_transform.py b/lightly/transforms/ssley_transform.py new file mode 100644 index 000000000..d5dbbc6ac --- /dev/null +++ b/lightly/transforms/ssley_transform.py @@ -0,0 +1,117 @@ +from typing import Dict, List, Optional, Tuple, Union + +from lightly.transforms.vicreg_transform import VICRegViewTransform + +from lightly.transforms.multi_view_transform import MultiViewTransform +from lightly.transforms.utils import IMAGENET_NORMALIZE + + +class SSLEYTransform(MultiViewTransform): + """Implements the transformations for SSL-EY. + + Input to this transform: + PIL Image or Tensor. + + Output of this transform: + List of Tensor of length 2. + + Applies the following augmentations by default: + - Random resized crop + - Random horizontal flip + - Color jitter + - Random gray scale + - Random solarization + - Gaussian blur + - ImageNet normalization + + Similar to SimCLR transform but with extra solarization. + + Attributes: + input_size: + Size of the input image in pixels. + cj_prob: + Probability that color jitter is applied. + cj_strength: + Strength of the color jitter. `cj_bright`, `cj_contrast`, `cj_sat`, and + `cj_hue` are multiplied by this value. + cj_bright: + How much to jitter brightness. + cj_contrast: + How much to jitter constrast. + cj_sat: + How much to jitter saturation. + cj_hue: + How much to jitter hue. + min_scale: + Minimum size of the randomized crop relative to the input_size. + random_gray_scale: + Probability of conversion to grayscale. + solarize_prob: + Probability of solarization. + gaussian_blur: + Probability of Gaussian blur. + kernel_size: + Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. + Used to calculate sigma of gaussian blur with kernel_size * input_size. + sigmas: + Tuple of min and max value from which the std of the gaussian kernel is sampled. + Is ignored if `kernel_size` is set. + vf_prob: + Probability that vertical flip is applied. + hf_prob: + Probability that horizontal flip is applied. + rr_prob: + Probability that random rotation is applied. + rr_degrees: + Range of degrees to select from for random rotation. If rr_degrees is None, + images are rotated by 90 degrees. If rr_degrees is a (min, max) tuple, + images are rotated by a random angle in [min, max]. If rr_degrees is a + single number, images are rotated by a random angle in + [-rr_degrees, +rr_degrees]. All rotations are counter-clockwise. + normalize: + Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. + + """ + + def __init__( + self, + input_size: int = 224, + cj_prob: float = 0.8, + cj_strength: float = 0.5, + cj_bright: float = 0.8, + cj_contrast: float = 0.8, + cj_sat: float = 0.4, + cj_hue: float = 0.2, + min_scale: float = 0.08, + random_gray_scale: float = 0.2, + solarize_prob: float = 0.1, + gaussian_blur: float = 0.5, + kernel_size: Optional[float] = None, + sigmas: Tuple[float, float] = (0.1, 2), + vf_prob: float = 0.0, + hf_prob: float = 0.5, + rr_prob: float = 0.0, + rr_degrees: Union[None, float, Tuple[float, float]] = None, + normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE, + ): + view_transform = VICRegViewTransform( + input_size=input_size, + cj_prob=cj_prob, + cj_strength=cj_strength, + cj_bright=cj_bright, + cj_contrast=cj_contrast, + cj_sat=cj_sat, + cj_hue=cj_hue, + min_scale=min_scale, + random_gray_scale=random_gray_scale, + solarize_prob=solarize_prob, + gaussian_blur=gaussian_blur, + kernel_size=kernel_size, + sigmas=sigmas, + vf_prob=vf_prob, + hf_prob=hf_prob, + rr_prob=rr_prob, + rr_degrees=rr_degrees, + normalize=normalize, + ) + super().__init__(transforms=[view_transform, view_transform]) \ No newline at end of file From 551d71bdd6603794bc80e69cdbe1241a3e3652ae Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:03:35 +0000 Subject: [PATCH 11/25] Format --- examples/pytorch/ssley.py | 1 + lightly/transforms/ssley_transform.py | 5 ++--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/pytorch/ssley.py b/examples/pytorch/ssley.py index 878737b56..8a5bfbfdc 100644 --- a/examples/pytorch/ssley.py +++ b/examples/pytorch/ssley.py @@ -1,6 +1,7 @@ import torch import torchvision from torch import nn + from lightly.loss.ssley_loss import SSLEYLoss from lightly.models.modules.heads import VICRegProjectionHead from lightly.transforms.vicreg_transform import VICRegTransform diff --git a/lightly/transforms/ssley_transform.py b/lightly/transforms/ssley_transform.py index d5dbbc6ac..00df33b3e 100644 --- a/lightly/transforms/ssley_transform.py +++ b/lightly/transforms/ssley_transform.py @@ -1,9 +1,8 @@ from typing import Dict, List, Optional, Tuple, Union -from lightly.transforms.vicreg_transform import VICRegViewTransform - from lightly.transforms.multi_view_transform import MultiViewTransform from lightly.transforms.utils import IMAGENET_NORMALIZE +from lightly.transforms.vicreg_transform import VICRegViewTransform class SSLEYTransform(MultiViewTransform): @@ -114,4 +113,4 @@ def __init__( rr_degrees=rr_degrees, normalize=normalize, ) - super().__init__(transforms=[view_transform, view_transform]) \ No newline at end of file + super().__init__(transforms=[view_transform, view_transform]) From bf3bab63ca3f30aa2eedd3c2a8b709911372b5ff Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:14:02 +0000 Subject: [PATCH 12/25] Cleanup types --- lightly/loss/ssley_loss.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py index b6023bd13..70ed0d5de 100644 --- a/lightly/loss/ssley_loss.py +++ b/lightly/loss/ssley_loss.py @@ -1,10 +1,12 @@ import torch import torch.distributed as dist +from torch import Tensor +from torch.nn import Module from lightly.utils.dist import gather -class SSLEYLoss(torch.nn.Module): +class SSLEYLoss(Module): """Implementation of the SSL-EY loss [0]. - [0] Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2022, https://arxiv.org/abs/2310.01012 @@ -35,9 +37,9 @@ class SSLEYLoss(torch.nn.Module): def __init__( self, gather_distributed: bool = False, - eps=0.0001, + eps: float = 0.0001, ): - super(SSLEYLoss, self).__init__() + super().__init__() if gather_distributed and not dist.is_available(): raise ValueError( "gather_distributed is True but torch.distributed is not available. " @@ -48,7 +50,7 @@ def __init__( self.gather_distributed = gather_distributed self.eps = eps - def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor: + def forward(self, z_a: Tensor, z_b: Tensor) -> Tensor: """Returns SSL-EY loss. Args: From 1c86d1b90e31e39012d3a5268f49c47d6e7852f3 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:14:12 +0000 Subject: [PATCH 13/25] Cleanup tests --- lightly/loss/ssley_loss.py | 15 +++++++++------ tests/loss/test_SSLEYLoss.py | 27 +++++++++++++-------------- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py index 70ed0d5de..02bcd0a8b 100644 --- a/lightly/loss/ssley_loss.py +++ b/lightly/loss/ssley_loss.py @@ -59,12 +59,15 @@ def forward(self, z_a: Tensor, z_b: Tensor) -> Tensor: z_b: Tensor with shape (batch_size, ..., dim). """ - assert ( - z_a.shape[0] > 1 and z_b.shape[0] > 1 - ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}" - assert ( - z_a.shape == z_b.shape - ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}." + if z_a.shape[0] <= 1: + raise ValueError(f"z_a must have batch size > 1 but found {z_a.shape[0]}.") + if z_b.shape[0] <= 1: + raise ValueError(f"z_b must have batch size > 1 but found {z_b.shape[0]}.") + if z_a.shape != z_b.shape: + raise ValueError( + f"z_a and z_b must have same shape but found {z_a.shape} and " + f"{z_b.shape}." + ) # gather all batches if self.gather_distributed and dist.is_initialized(): world_size = dist.get_world_size() diff --git a/tests/loss/test_SSLEYLoss.py b/tests/loss/test_SSLEYLoss.py index 0ee7b1e2b..838b0019a 100644 --- a/tests/loss/test_SSLEYLoss.py +++ b/tests/loss/test_SSLEYLoss.py @@ -9,12 +9,12 @@ class TestSSLEYLoss: - def test__gather_distributed(self, mocker: MockerFixture) -> None: + def test__init_gather_distributed(self, mocker: MockerFixture) -> None: mock_is_available = mocker.patch.object(dist, "is_available", return_value=True) SSLEYLoss(gather_distributed=True) mock_is_available.assert_called_once() - def test__gather_distributed_dist_not_available( + def test__init_gather_distributed_dist_not_available( self, mocker: MockerFixture ) -> None: mock_is_available = mocker.patch.object( @@ -24,10 +24,7 @@ def test__gather_distributed_dist_not_available( SSLEYLoss(gather_distributed=True) mock_is_available.assert_called_once() - -class TestSSLEYLossUnitTest(unittest.TestCase): - # Old tests in unittest style, please add new tests to TestSSLEYLoss using pytest. - def test_forward_pass(self): + def test_forward(self) -> None: loss = SSLEYLoss() for bsz in range(2, 4): x0 = torch.randn((bsz, 32)) @@ -36,10 +33,10 @@ def test_forward_pass(self): # symmetry l1 = loss(x0, x1) l2 = loss(x1, x0) - self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0) + assert (l1 - l2).pow(2).item() == pytest.approx(0.0) @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available") - def test_forward_pass_cuda(self): + def test_forward__cuda(self): loss = SSLEYLoss() for bsz in range(2, 4): x0 = torch.randn((bsz, 32)).cuda() @@ -48,18 +45,20 @@ def test_forward_pass_cuda(self): # symmetry l1 = loss(x0, x1) l2 = loss(x1, x0) - self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0) + assert (l1 - l2).pow(2).item() == pytest.approx(0.0) - def test_forward_pass__error_batch_size_1(self): + def test_forward__error_batch_size_1(self): loss = SSLEYLoss() x0 = torch.randn((1, 32)) - x1 = torch.randn((1, 32)) - with self.assertRaises(AssertionError): + x1 = torch.randn((2, 32)) + with pytest.raises(ValueError): loss(x0, x1) + with pytest.raises(ValueError): + loss(x1, x0) - def test_forward_pass__error_different_shapes(self): + def test_forward__error_different_shapes(self): loss = SSLEYLoss() x0 = torch.randn((2, 32)) x1 = torch.randn((2, 16)) - with self.assertRaises(AssertionError): + with pytest.raises(ValueError): loss(x0, x1) From 6f82ba38942f0e4dd36525e627ede75bca820f36 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:16:26 +0000 Subject: [PATCH 14/25] Use SSLEYTransform --- examples/pytorch_lightning_distributed/ssley.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/pytorch_lightning_distributed/ssley.py b/examples/pytorch_lightning_distributed/ssley.py index 3f034adcf..96cb68936 100644 --- a/examples/pytorch_lightning_distributed/ssley.py +++ b/examples/pytorch_lightning_distributed/ssley.py @@ -11,7 +11,7 @@ ## The projection head is the same as the Barlow Twins one from lightly.models.modules.heads import VICRegProjectionHead -from lightly.transforms.vicreg_transform import VICRegTransform +from lightly.transforms.ssley_transform import SSLEYTransform class SSLEY(pl.LightningModule): @@ -49,7 +49,7 @@ def configure_optimizers(self): model = SSLEY() -transform = VICRegTransform(input_size=32) +transform = SSLEYTransform(input_size=32) dataset = torchvision.datasets.CIFAR10( "datasets/cifar10", download=True, transform=transform ) From b8f49205ce1e2d208d434b7acbee24ef02b03159 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:21:12 +0000 Subject: [PATCH 15/25] Fix types --- tests/loss/test_SSLEYLoss.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/loss/test_SSLEYLoss.py b/tests/loss/test_SSLEYLoss.py index 838b0019a..02b317bc4 100644 --- a/tests/loss/test_SSLEYLoss.py +++ b/tests/loss/test_SSLEYLoss.py @@ -36,7 +36,7 @@ def test_forward(self) -> None: assert (l1 - l2).pow(2).item() == pytest.approx(0.0) @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available") - def test_forward__cuda(self): + def test_forward__cuda(self) -> None: loss = SSLEYLoss() for bsz in range(2, 4): x0 = torch.randn((bsz, 32)).cuda() @@ -47,7 +47,7 @@ def test_forward__cuda(self): l2 = loss(x1, x0) assert (l1 - l2).pow(2).item() == pytest.approx(0.0) - def test_forward__error_batch_size_1(self): + def test_forward__error_batch_size_1(self) -> None: loss = SSLEYLoss() x0 = torch.randn((1, 32)) x1 = torch.randn((2, 32)) @@ -56,7 +56,7 @@ def test_forward__error_batch_size_1(self): with pytest.raises(ValueError): loss(x1, x0) - def test_forward__error_different_shapes(self): + def test_forward__error_different_shapes(self) -> None: loss = SSLEYLoss() x0 = torch.randn((2, 32)) x1 = torch.randn((2, 16)) From 66d278761707880e22469465baf99890804d34b4 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:21:20 +0000 Subject: [PATCH 16/25] Add SSLEYTransform test --- lightly/transforms/__init__.py | 1 + tests/transforms/test_ssley_transform.py | 13 +++++++++++++ 2 files changed, 14 insertions(+) create mode 100644 tests/transforms/test_ssley_transform.py diff --git a/lightly/transforms/__init__.py b/lightly/transforms/__init__.py index 5c3ddbba6..30a32ace9 100644 --- a/lightly/transforms/__init__.py +++ b/lightly/transforms/__init__.py @@ -31,6 +31,7 @@ from lightly.transforms.simsiam_transform import SimSiamTransform, SimSiamViewTransform from lightly.transforms.smog_transform import SMoGTransform, SmoGViewTransform from lightly.transforms.solarize import RandomSolarization +from lightly.transforms.ssley_transform import SSLEYTransform from lightly.transforms.swav_transform import SwaVTransform, SwaVViewTransform from lightly.transforms.vicreg_transform import VICRegTransform, VICRegViewTransform from lightly.transforms.vicregl_transform import VICRegLTransform, VICRegLViewTransform diff --git a/tests/transforms/test_ssley_transform.py b/tests/transforms/test_ssley_transform.py new file mode 100644 index 000000000..3ff49672d --- /dev/null +++ b/tests/transforms/test_ssley_transform.py @@ -0,0 +1,13 @@ +from PIL import Image + +from lightly.transforms import SSLEYTransform + + +class TestSSLEYTransform: + def test__call__(self) -> None: + transform = SSLEYTransform(input_size=32) + + sample = Image.new("RGB", (100, 100)) + output = transform(sample) + assert len(output) == 2 + assert all(out.shape == (3, 32, 32) for out in output) From 62cf5b4107082ea159cecad1dbb6f6350a432d43 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:28:42 +0000 Subject: [PATCH 17/25] Add SSLEYProjectionHead --- lightly/models/modules/heads.py | 24 ++++++++++++++++++++++++ tests/models/test_ProjectionHeads.py | 6 ++++-- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/lightly/models/modules/heads.py b/lightly/models/modules/heads.py index 541b49a29..a2b789866 100644 --- a/lightly/models/modules/heads.py +++ b/lightly/models/modules/heads.py @@ -750,3 +750,27 @@ def __init__( (hidden_dim, output_dim, None, None), ] ) + + +class SSLEYProjectionHead(VICRegProjectionHead): + """Projection head used for SSLEY [0]. + + Same as VICReg [1] projection head but uses only 2 layers by default. + + - [0]: 2023, Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, https://arxiv.org/abs/2310.01012 + - [1]: 2022, VICRegL, https://arxiv.org/abs/2210.01571 + """ + + def __init__( + self, + input_dim: int = 2048, + hidden_dim: int = 8192, + output_dim: int = 8192, + num_layers: int = 2, + ): + super().__init__( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + num_layers=num_layers, + ) diff --git a/tests/models/test_ProjectionHeads.py b/tests/models/test_ProjectionHeads.py index f98b33143..131028d1a 100644 --- a/tests/models/test_ProjectionHeads.py +++ b/tests/models/test_ProjectionHeads.py @@ -2,7 +2,7 @@ import torch -import lightly +from lightly.loss import DINOLoss from lightly.models.modules.heads import ( BarlowTwinsProjectionHead, BYOLPredictionHead, @@ -15,6 +15,7 @@ SimCLRProjectionHead, SimSiamPredictionHead, SimSiamProjectionHead, + SSLEYProjectionHead, SwaVProjectionHead, SwaVPrototypes, TiCoProjectionHead, @@ -45,6 +46,7 @@ def setUp(self): SimCLRProjectionHead, SimSiamProjectionHead, SimSiamPredictionHead, + SSLEYProjectionHead, SwaVProjectionHead, TiCoProjectionHead, VicRegLLocalProjectionHead, @@ -192,7 +194,7 @@ def test_dino_projection_head_freeze_last_layer(self, seed=0): norm_last_layer=norm_last_layer, ) optimizer = torch.optim.SGD(head.parameters(), lr=1) - criterion = lightly.loss.DINOLoss(output_dim=4) + criterion = DINOLoss(output_dim=4) # Store initial weights of last layer initial_data = [ param.data.detach().clone() From 6b9bc355542bd465f7e68e5cb81bd5dcfc8eb9bc Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:33:07 +0000 Subject: [PATCH 18/25] Update examples --- benchmarks/imagenet/resnet50/ssley.py | 8 ++++---- examples/pytorch/ssley.py | 8 ++++---- examples/pytorch_lightning/ssley.py | 10 ++++------ examples/pytorch_lightning_distributed/ssley.py | 8 +++----- 4 files changed, 15 insertions(+), 19 deletions(-) diff --git a/benchmarks/imagenet/resnet50/ssley.py b/benchmarks/imagenet/resnet50/ssley.py index b789b0c5d..83d6b0237 100644 --- a/benchmarks/imagenet/resnet50/ssley.py +++ b/benchmarks/imagenet/resnet50/ssley.py @@ -7,9 +7,9 @@ from torchvision.models import resnet50 from lightly.loss.ssley_loss import SSLEYLoss -from lightly.models.modules.heads import VICRegProjectionHead +from lightly.models.modules.heads import SSLEYProjectionHead from lightly.models.utils import get_weight_decay_parameters -from lightly.transforms.vicreg_transform import VICRegTransform +from lightly.transforms import SSLEYTransform from lightly.utils.benchmarking import OnlineLinearClassifier from lightly.utils.lars import LARS from lightly.utils.scheduler import CosineWarmupScheduler @@ -24,7 +24,7 @@ def __init__(self, batch_size_per_device: int, num_classes: int) -> None: resnet = resnet50() resnet.fc = Identity() # Ignore classification head self.backbone = resnet - self.projection_head = VICRegProjectionHead(num_layers=2) + self.projection_head = SSLEYProjectionHead() self.criterion = SSLEYLoss() self.online_classifier = OnlineLinearClassifier(num_classes=num_classes) @@ -108,7 +108,7 @@ def configure_optimizers(self): # SSLEY transform -transform = VICRegTransform() +transform = SSLEYTransform() def _get_base_learning_rate(global_batch_size: int) -> float: diff --git a/examples/pytorch/ssley.py b/examples/pytorch/ssley.py index 8a5bfbfdc..964e34757 100644 --- a/examples/pytorch/ssley.py +++ b/examples/pytorch/ssley.py @@ -3,15 +3,15 @@ from torch import nn from lightly.loss.ssley_loss import SSLEYLoss -from lightly.models.modules.heads import VICRegProjectionHead -from lightly.transforms.vicreg_transform import VICRegTransform +from lightly.models.modules.heads import SSLEYProjectionHead +from lightly.transforms import SSLEYTransform class SSLEY(nn.Module): def __init__(self, backbone): super().__init__() self.backbone = backbone - self.projection_head = VICRegProjectionHead( + self.projection_head = SSLEYProjectionHead( input_dim=512, hidden_dim=2048, output_dim=2048, @@ -31,7 +31,7 @@ def forward(self, x): device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) -transform = VICRegTransform(input_size=32) +transform = SSLEYTransform(input_size=32) dataset = torchvision.datasets.CIFAR10( "datasets/cifar10", download=True, transform=transform ) diff --git a/examples/pytorch_lightning/ssley.py b/examples/pytorch_lightning/ssley.py index 6cc17de96..fa520ee34 100644 --- a/examples/pytorch_lightning/ssley.py +++ b/examples/pytorch_lightning/ssley.py @@ -8,10 +8,8 @@ from torch import nn from lightly.loss.ssley_loss import SSLEYLoss - -## The projection head is the same as the Barlow Twins one -from lightly.models.modules.heads import VICRegProjectionHead -from lightly.transforms.vicreg_transform import VICRegTransform +from lightly.models.modules.heads import SSLEYProjectionHead +from lightly.transforms import SSLEYTransform class SSLEY(pl.LightningModule): @@ -19,7 +17,7 @@ def __init__(self): super().__init__() resnet = torchvision.models.resnet18() self.backbone = nn.Sequential(*list(resnet.children())[:-1]) - self.projection_head = VICRegProjectionHead( + self.projection_head = SSLEYProjectionHead( input_dim=512, hidden_dim=2048, output_dim=2048, @@ -46,7 +44,7 @@ def configure_optimizers(self): model = SSLEY() -transform = VICRegTransform(input_size=32) +transform = SSLEYTransform(input_size=32) dataset = torchvision.datasets.CIFAR10( "datasets/cifar10", download=True, transform=transform ) diff --git a/examples/pytorch_lightning_distributed/ssley.py b/examples/pytorch_lightning_distributed/ssley.py index 96cb68936..7d0ae6e93 100644 --- a/examples/pytorch_lightning_distributed/ssley.py +++ b/examples/pytorch_lightning_distributed/ssley.py @@ -8,10 +8,8 @@ from torch import nn from lightly.loss import SSLEYLoss - -## The projection head is the same as the Barlow Twins one -from lightly.models.modules.heads import VICRegProjectionHead -from lightly.transforms.ssley_transform import SSLEYTransform +from lightly.models.modules.heads import SSLEYProjectionHead +from lightly.transforms import SSLEYTransform class SSLEY(pl.LightningModule): @@ -19,7 +17,7 @@ def __init__(self): super().__init__() resnet = torchvision.models.resnet18() self.backbone = nn.Sequential(*list(resnet.children())[:-1]) - self.projection_head = VICRegProjectionHead( + self.projection_head = SSLEYProjectionHead( input_dim=512, hidden_dim=2048, output_dim=2048, From b4bfa87543f7d4c99b2d44a2b476b0a354eafdd6 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:34:59 +0000 Subject: [PATCH 19/25] Fix citation --- lightly/loss/ssley_loss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py index 02bcd0a8b..779ff120d 100644 --- a/lightly/loss/ssley_loss.py +++ b/lightly/loss/ssley_loss.py @@ -9,7 +9,7 @@ class SSLEYLoss(Module): """Implementation of the SSL-EY loss [0]. - - [0] Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2022, https://arxiv.org/abs/2310.01012 + - [0]: Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2023, https://arxiv.org/abs/2310.01012 Attributes: gather_distributed: From 1f4c5a899d414ad5627f735cd18705af3779dd38 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 13:55:54 +0000 Subject: [PATCH 20/25] Fix typings in dist.py --- lightly/utils/dist.py | 42 ++++++++++++++++++++++++++++++------------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/lightly/utils/dist.py b/lightly/utils/dist.py index 2c20750c6..e3fee2a99 100644 --- a/lightly/utils/dist.py +++ b/lightly/utils/dist.py @@ -1,10 +1,13 @@ -from typing import Optional, Tuple +from typing import Any, Callable, Literal, Optional, Tuple, TypeVar, Union +from torch.autograd import Function +from torch.autograd.function import FunctionCtx import torch import torch.distributed as dist +from torch import Tensor -class GatherLayer(torch.autograd.Function): +class GatherLayer(Function): """Gather tensors from all processes, supporting backward propagation. This code was taken and adapted from here: @@ -13,15 +16,16 @@ class GatherLayer(torch.autograd.Function): """ @staticmethod - def forward(ctx, input: torch.Tensor) -> Tuple[torch.Tensor, ...]: + def forward(ctx: FunctionCtx, input: Tensor) -> Tuple[Tensor, ...]: ctx.save_for_backward(input) output = [torch.empty_like(input) for _ in range(dist.get_world_size())] dist.all_gather(output, input) return tuple(output) @staticmethod - def backward(ctx, *grads: torch.Tensor) -> torch.Tensor: - (input,) = ctx.saved_tensors + def backward(ctx: FunctionCtx, *grads: Tensor) -> Tensor: + # Type ignore because FunctionCtx is not fully typed. + (input,) = ctx.saved_tensors # type: ignore grad_out = torch.empty_like(input) grad_out[:] = grads[dist.get_rank()] return grad_out @@ -37,12 +41,13 @@ def world_size() -> int: return dist.get_world_size() if dist.is_initialized() else 1 -def gather(input: torch.Tensor) -> Tuple[torch.Tensor]: +def gather(input: Tensor) -> Tuple[Tensor]: """Gathers this tensor from all processes. Supports backprop.""" - return GatherLayer.apply(input) + # Type ignore is required because Function.apply is untyped. + return GatherLayer.apply(input) # type: ignore -def eye_rank(n: int, device: Optional[torch.device] = None) -> torch.Tensor: +def eye_rank(n: int, device: Optional[torch.device] = None) -> Tensor: """Returns an (n, n * world_size) zero matrix with the diagonal for the rank of this process set to 1. @@ -70,7 +75,12 @@ def eye_rank(n: int, device: Optional[torch.device] = None) -> torch.Tensor: return diag_mask -def rank_zero_only(fn): +_T = TypeVar("_T") + + +# TODO(Guarin, 01/2024): Refine typings for callable with ParamSpec once we drop support +# for Python <=3.9. +def rank_zero_only(fn: Callable[..., _T]) -> Callable[..., Union[_T, None]]: """Decorator that only runs the function on the process with rank 0. Example: @@ -82,14 +92,22 @@ def rank_zero_only(fn): """ - def wrapped(*args, **kwargs): + def wrapped(*args: Any, **kwargs: Any) -> Union[_T, None]: if rank() == 0: return fn(*args, **kwargs) + return None return wrapped +# Type ignore required because 'file' has Any type. @rank_zero_only -def print_rank_zero(*args, **kwargs) -> None: +def print_rank_zero( # type: ignore[misc] + *values: object, + sep: Union[str, None] = " ", + end: Union[str, None] = "\n", + file: Any = None, + flush: Literal[False] = False +) -> None: """Equivalent to print, but only runs on the process with rank 0.""" - print(*args, **kwargs) + print(*values, sep=sep, end=end, file=file, flush=flush) From 7e3720f246234572faf7ea093d228d995684f1c8 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 14:11:57 +0000 Subject: [PATCH 21/25] Format --- lightly/utils/dist.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lightly/utils/dist.py b/lightly/utils/dist.py index e3fee2a99..e56b43a3c 100644 --- a/lightly/utils/dist.py +++ b/lightly/utils/dist.py @@ -1,10 +1,10 @@ from typing import Any, Callable, Literal, Optional, Tuple, TypeVar, Union -from torch.autograd import Function -from torch.autograd.function import FunctionCtx import torch import torch.distributed as dist from torch import Tensor +from torch.autograd import Function +from torch.autograd.function import FunctionCtx class GatherLayer(Function): @@ -107,7 +107,7 @@ def print_rank_zero( # type: ignore[misc] sep: Union[str, None] = " ", end: Union[str, None] = "\n", file: Any = None, - flush: Literal[False] = False + flush: bool = False ) -> None: """Equivalent to print, but only runs on the process with rank 0.""" print(*values, sep=sep, end=end, file=file, flush=flush) From b9d0ed9e3ee5fcaa5c41035f0af3dd1f169a1432 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 14:13:25 +0000 Subject: [PATCH 22/25] Remove unused import --- lightly/utils/dist.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightly/utils/dist.py b/lightly/utils/dist.py index e56b43a3c..7c8c7f773 100644 --- a/lightly/utils/dist.py +++ b/lightly/utils/dist.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Literal, Optional, Tuple, TypeVar, Union +from typing import Any, Callable, Optional, Tuple, TypeVar, Union import torch import torch.distributed as dist From 9261583aece340e8b1ee2f06c68f017ce8da468c Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 14:24:54 +0000 Subject: [PATCH 23/25] Use any instead of FunctionCtx --- lightly/utils/dist.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/lightly/utils/dist.py b/lightly/utils/dist.py index 7c8c7f773..f44df6b9b 100644 --- a/lightly/utils/dist.py +++ b/lightly/utils/dist.py @@ -4,7 +4,6 @@ import torch.distributed as dist from torch import Tensor from torch.autograd import Function -from torch.autograd.function import FunctionCtx class GatherLayer(Function): @@ -15,17 +14,18 @@ class GatherLayer(Function): """ + # Type ignore is required because superclass uses Any type for ctx. @staticmethod - def forward(ctx: FunctionCtx, input: Tensor) -> Tuple[Tensor, ...]: + def forward(ctx: Any, input: Tensor) -> Tuple[Tensor, ...]: # type: ignore[misc] ctx.save_for_backward(input) output = [torch.empty_like(input) for _ in range(dist.get_world_size())] dist.all_gather(output, input) return tuple(output) + # Type ignore is required because superclass uses Any type for ctx. @staticmethod - def backward(ctx: FunctionCtx, *grads: Tensor) -> Tensor: - # Type ignore because FunctionCtx is not fully typed. - (input,) = ctx.saved_tensors # type: ignore + def backward(ctx: Any, *grads: Tensor) -> Tensor: # type: ignore[misc] + (input,) = ctx.saved_tensors grad_out = torch.empty_like(input) grad_out[:] = grads[dist.get_rank()] return grad_out From e8e30547893c2979e648108284c17d8ca8ec2995 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 14:35:05 +0000 Subject: [PATCH 24/25] Add type ignore override --- lightly/utils/dist.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/lightly/utils/dist.py b/lightly/utils/dist.py index f44df6b9b..88407556d 100644 --- a/lightly/utils/dist.py +++ b/lightly/utils/dist.py @@ -14,9 +14,11 @@ class GatherLayer(Function): """ - # Type ignore is required because superclass uses Any type for ctx. + # Type ignore misc is required because the superclass uses Any type for ctx. + # Type ignore override is required because the superclass has a different signature + # for forward. @staticmethod - def forward(ctx: Any, input: Tensor) -> Tuple[Tensor, ...]: # type: ignore[misc] + def forward(ctx: Any, input: Tensor) -> Tuple[Tensor, ...]: # type: ignore[misc, override] ctx.save_for_backward(input) output = [torch.empty_like(input) for _ in range(dist.get_world_size())] dist.all_gather(output, input) From f88efd0ac95801b8e6106e9b44489ea8fcc66cb6 Mon Sep 17 00:00:00 2001 From: guarin Date: Fri, 5 Jan 2024 14:42:13 +0000 Subject: [PATCH 25/25] Show SSL-EY in docs --- docs/source/examples/models.rst | 1 + docs/source/lightly.transforms.rst | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/docs/source/examples/models.rst b/docs/source/examples/models.rst index ec9fe28e5..9d2b97710 100644 --- a/docs/source/examples/models.rst +++ b/docs/source/examples/models.rst @@ -24,6 +24,7 @@ for PyTorch and PyTorch Lightning to give you a headstart when implementing your simmim.rst simsiam.rst smog.rst + ssley.rst swav.rst tico.rst vicreg.rst diff --git a/docs/source/lightly.transforms.rst b/docs/source/lightly.transforms.rst index 7b4573fd1..d47c8da32 100644 --- a/docs/source/lightly.transforms.rst +++ b/docs/source/lightly.transforms.rst @@ -71,6 +71,10 @@ lightly.transforms :members: :special-members: __call__ +.. automodule:: lightly.transforms.ssley_transform + :members: + :special-members: __call__ + .. automodule:: lightly.transforms.swav_transform :members: :special-members: __call__