diff --git a/README.md b/README.md index 4f61264cd..0521808a3 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) @@ -287,6 +288,7 @@ See the [benchmarking scripts](./benchmarks/imagenet/resnet50/) for details. | 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) | diff --git a/benchmarks/imagenet/resnet50/main.py b/benchmarks/imagenet/resnet50/main.py index 34ce68f12..dab7b7851 100644 --- a/benchmarks/imagenet/resnet50/main.py +++ b/benchmarks/imagenet/resnet50/main.py @@ -13,6 +13,7 @@ import linear_eval import mocov2 import simclr +import ssley import swav import torch import vicreg @@ -60,6 +61,7 @@ "dino": {"model": dino.DINO, "transform": dino.transform}, "mocov2": {"model": mocov2.MoCoV2, "transform": mocov2.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/benchmarks/imagenet/resnet50/ssley.py b/benchmarks/imagenet/resnet50/ssley.py new file mode 100644 index 000000000..83d6b0237 --- /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 SSLEYProjectionHead +from lightly.models.utils import get_weight_decay_parameters +from lightly.transforms import SSLEYTransform +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 = SSLEYProjectionHead() + 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 = SSLEYTransform() + + +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/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/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 diff --git a/docs/source/getting_started/benchmarks/imagenette_benchmark.py b/docs/source/getting_started/benchmarks/imagenette_benchmark.py index ee9cc0171..e94fa3d5d 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, @@ -266,6 +269,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, @@ -1166,6 +1170,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) @@ -1411,6 +1451,7 @@ def configure_optimizers(self): SimCLRModel, # SimMIMModel, # disabled by default because SimMIM uses larger images with size 224 SimSiamModel, + SSLEYModel, SwaVModel, SwaVQueueModel, SMoGModel, 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/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__ diff --git a/examples/pytorch/ssley.py b/examples/pytorch/ssley.py new file mode 100644 index 000000000..964e34757 --- /dev/null +++ b/examples/pytorch/ssley.py @@ -0,0 +1,66 @@ +import torch +import torchvision +from torch import nn + +from lightly.loss.ssley_loss import SSLEYLoss +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 = SSLEYProjectionHead( + 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 = SSLEYTransform(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..fa520ee34 --- /dev/null +++ b/examples/pytorch_lightning/ssley.py @@ -0,0 +1,65 @@ +# 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 +from lightly.models.modules.heads import SSLEYProjectionHead +from lightly.transforms import SSLEYTransform + + +class SSLEY(pl.LightningModule): + def __init__(self): + super().__init__() + resnet = torchvision.models.resnet18() + self.backbone = nn.Sequential(*list(resnet.children())[:-1]) + self.projection_head = SSLEYProjectionHead( + 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 = SSLEYTransform(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..7d0ae6e93 --- /dev/null +++ b/examples/pytorch_lightning_distributed/ssley.py @@ -0,0 +1,75 @@ +# 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 +from lightly.models.modules.heads import SSLEYProjectionHead +from lightly.transforms import SSLEYTransform + + +class SSLEY(pl.LightningModule): + def __init__(self): + super().__init__() + resnet = torchvision.models.resnet18() + self.backbone = nn.Sequential(*list(resnet.children())[:-1]) + self.projection_head = SSLEYProjectionHead( + 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 = SSLEYTransform(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/__init__.py b/lightly/loss/__init__.py index 85e94e042..397a3b345 100644 --- a/lightly/loss/__init__.py +++ b/lightly/loss/__init__.py @@ -10,6 +10,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 diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py new file mode 100644 index 000000000..779ff120d --- /dev/null +++ b/lightly/loss/ssley_loss.py @@ -0,0 +1,86 @@ +import torch +import torch.distributed as dist +from torch import Tensor +from torch.nn import Module + +from lightly.utils.dist import gather + + +class SSLEYLoss(Module): + """Implementation of the SSL-EY loss [0]. + + - [0]: Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2023, 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: float = 0.0001, + ): + super().__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.gather_distributed = gather_distributed + self.eps = eps + + def forward(self, z_a: Tensor, z_b: Tensor) -> Tensor: + """Returns SSL-EY loss. + + Args: + z_a: + Tensor with shape (batch_size, ..., dim). + z_b: + Tensor with shape (batch_size, ..., dim). + """ + 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() + if world_size > 1: + z_a = torch.cat(gather(z_a), dim=0) + z_b = torch.cat(gather(z_b), dim=0) + + 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) + + loss = -2 * torch.trace(C) + torch.trace(V @ V) + + return loss 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/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/lightly/transforms/ssley_transform.py b/lightly/transforms/ssley_transform.py new file mode 100644 index 000000000..00df33b3e --- /dev/null +++ b/lightly/transforms/ssley_transform.py @@ -0,0 +1,116 @@ +from typing import Dict, List, Optional, Tuple, Union + +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): + """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]) diff --git a/lightly/utils/dist.py b/lightly/utils/dist.py index 2c20750c6..88407556d 100644 --- a/lightly/utils/dist.py +++ b/lightly/utils/dist.py @@ -1,10 +1,12 @@ -from typing import Optional, Tuple +from typing import Any, Callable, Optional, Tuple, TypeVar, Union import torch import torch.distributed as dist +from torch import Tensor +from torch.autograd import Function -class GatherLayer(torch.autograd.Function): +class GatherLayer(Function): """Gather tensors from all processes, supporting backward propagation. This code was taken and adapted from here: @@ -12,15 +14,19 @@ class GatherLayer(torch.autograd.Function): """ + # 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, input: torch.Tensor) -> Tuple[torch.Tensor, ...]: + 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) return tuple(output) + # Type ignore is required because superclass uses Any type for ctx. @staticmethod - def backward(ctx, *grads: torch.Tensor) -> torch.Tensor: + 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()] @@ -37,12 +43,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 +77,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 +94,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: bool = 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) diff --git a/tests/loss/test_SSLEYLoss.py b/tests/loss/test_SSLEYLoss.py new file mode 100644 index 000000000..02b317bc4 --- /dev/null +++ b/tests/loss/test_SSLEYLoss.py @@ -0,0 +1,64 @@ +import unittest + +import pytest +import torch +from pytest_mock import MockerFixture +from torch import distributed as dist + +from lightly.loss import SSLEYLoss + + +class TestSSLEYLoss: + 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__init_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() + + def test_forward(self) -> None: + 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) + assert (l1 - l2).pow(2).item() == pytest.approx(0.0) + + @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available") + def test_forward__cuda(self) -> None: + 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) + assert (l1 - l2).pow(2).item() == pytest.approx(0.0) + + def test_forward__error_batch_size_1(self) -> None: + loss = SSLEYLoss() + x0 = torch.randn((1, 32)) + x1 = torch.randn((2, 32)) + with pytest.raises(ValueError): + loss(x0, x1) + with pytest.raises(ValueError): + loss(x1, x0) + + def test_forward__error_different_shapes(self) -> None: + loss = SSLEYLoss() + x0 = torch.randn((2, 32)) + x1 = torch.randn((2, 16)) + with pytest.raises(ValueError): + loss(x0, x1) 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() 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)