Implements SSL-EY

docs/source/solo/losses/ssley.rst

@@ -0,0 +1,5 @@
+SSL-EY
+-------
+
+.. autofunction:: solo.losses.ssley.ssley_loss_func
+   :noindex:

docs/source/solo/methods/ssley.rst

@@ -0,0 +1,26 @@
+SSL-EY
+=======
+
+
+.. automethod:: solo.methods.ssley.SSLEY.__init__
+   :noindex:
+
+add_model_specific_args
+~~~~~~~~~~~~~~~~~~~~~~~
+.. automethod:: solo.methods.ssley.SSLEY.add_model_specific_args
+   :noindex:
+
+learnable_params
+~~~~~~~~~~~~~~~~
+.. autoattribute:: solo.methods.ssley.SSLEY.learnable_params
+   :noindex:
+
+forward
+~~~~~~~
+.. automethod:: solo.methods.ssley.SSLEY.forward
+   :noindex:
+
+training_step
+~~~~~~~~~~~~~
+.. automethod:: solo.methods.ssley.SSLEY.training_step
+   :noindex:

docs/source/start/available.rst

@@ -11,6 +11,7 @@ Methods available
 * `SwAV <https://arxiv.org/abs/2006.09882>`_
 * `VICReg <https://arxiv.org/abs/2105.04906>`_
 * `W-MSE <https://arxiv.org/abs/2007.06346>`_
+* `SSL-EYE <https://arxiv.org/abs/2310.01012>`_
 
 ************
 Extra flavor

scripts/linear/imagenet-100/ssley.yaml

@@ -0,0 +1,45 @@
+defaults:
+  - _self_
+  - wandb: private.yaml
+  - override hydra/hydra_logging: disabled
+  - override hydra/job_logging: disabled
+
+# disable hydra outputs
+hydra:
+  output_subdir: null
+  run:
+    dir: .
+
+name: "ssley-imagenet100-linear"
+pretrained_feature_extractor: None
+backbone:
+  name: "resnet18"
+pretrain_method: "ssley"
+data:
+  dataset: imagenet100
+  train_path: "./datasets/imagenet-100/train"
+  val_path: "./datasets/imagenet-100/val"
+  format: "dali"
+  num_workers: 4
+optimizer:
+  name: "sgd"
+  batch_size: 256
+  lr: 0.3
+  weight_decay: 0
+scheduler:
+  name: "step"
+  lr_decay_steps: [60, 80]
+checkpoint:
+  enabled: True
+  dir: "trained_models"
+  frequency: 1
+auto_resume:
+  enabled: True
+
+# overwrite PL stuff
+max_epochs: 100
+devices: [0]
+sync_batchnorm: True
+accelerator: "gpu"
+strategy: "ddp"
+precision: 16

scripts/pretrain/cifar/ssley.yaml

@@ -0,0 +1,80 @@
+defaults:
+  - _self_
+  - wandb: private.yaml
+  - override hydra/hydra_logging: disabled
+  - override hydra/job_logging: disabled
+
+# disable hydra outputs
+hydra:
+  output_subdir: null
+  run:
+    dir: .
+
+name: "ssley-cifar10" # change here for cifar100
+method: "ssley"
+backbone:
+  name: "resnet18"
+method_kwargs:
+  proj_hidden_dim: 2048
+  proj_output_dim: 2048
+data:
+  dataset: cifar10 # change here for cifar100
+  train_path: "./datasets"
+  val_path: "datasets/imagenet100/val"
+  format: "image_folder"
+  num_workers: 4
+augmentations:
+  - rrc:
+      enabled: True
+      crop_min_scale: 0.2
+      crop_max_scale: 1.0
+    color_jitter:
+      enabled: True
+      brightness: 0.4
+      contrast: 0.4
+      saturation: 0.2
+      hue: 0.1
+      prob: 0.8
+    grayscale:
+      enabled: True
+      prob: 0.2
+    gaussian_blur:
+      enabled: False
+      prob: 0.0
+    solarization:
+      enabled: True
+      prob: 0.1
+    equalization:
+      enabled: False
+      prob: 0.0
+    horizontal_flip:
+      enabled: True
+      prob: 0.5
+    crop_size: 32
+    num_crops: 2
+optimizer:
+  name: "lars"
+  batch_size: 256
+  lr: 0.3
+  classifier_lr: 0.1
+  weight_decay: 1e-4
+  kwargs:
+    clip_lr: True
+    eta: 0.02
+    exclude_bias_n_norm: True
+scheduler:
+  name: "warmup_cosine"
+checkpoint:
+  enabled: True
+  dir: "trained_models"
+  frequency: 1
+auto_resume:
+  enabled: True
+
+# overwrite PL stuff
+max_epochs: 1000
+devices: [0]
+sync_batchnorm: True
+accelerator: "gpu"
+strategy: "ddp"
+precision: 16-mixed

scripts/pretrain/imagenet-100/ssley.yaml

@@ -0,0 +1,81 @@
+defaults:
+  - _self_
+  - augmentations: ssley.yaml
+  - wandb: private.yaml
+  - override hydra/hydra_logging: disabled
+  - override hydra/job_logging: disabled
+
+# disable hydra outputs
+hydra:
+  output_subdir: null
+  run:
+    dir: .
+
+name: "ssley-imagenet100"
+method: "ssley"
+backbone:
+  name: "resnet18"
+method_kwargs:
+  proj_hidden_dim: 2048
+  proj_output_dim: 2048
+data:
+  dataset: imagenet100
+  train_path: "datasets/imagenet100/train"
+  val_path: "datasets/imagenet100/val"
+  format: "dali"
+  num_workers: 4
+augmentations:
+  - rrc:
+      enabled: True
+      crop_min_scale: 0.2
+      crop_max_scale: 1.0
+    color_jitter:
+      enabled: True
+      brightness: 0.4
+      contrast: 0.4
+      saturation: 0.2
+      hue: 0.1
+      prob: 0.8
+    grayscale:
+      enabled: True
+      prob: 0.2
+    gaussian_blur:
+      enabled: True
+      prob: 0.5
+    solarization:
+      enabled: True
+      prob: 0.1
+    equalization:
+      enabled: False
+      prob: 0.0
+    horizontal_flip:
+      enabled: True
+      prob: 0.5
+    crop_size: 224
+    num_crops: 2
+optimizer:
+  name: "lars"
+  batch_size: 128
+  lr: 0.3
+  classifier_lr: 0.1
+  weight_decay: 1e-4
+  kwargs:
+    clip_lr: True
+    eta: 0.02
+    exclude_bias_n_norm: True
+scheduler:
+  name: "warmup_cosine"
+checkpoint:
+  enabled: True
+  dir: "trained_models"
+  frequency: 1
+auto_resume:
+  enabled: True
+
+# overwrite PL stuff
+max_epochs: 400
+devices: [0, 1]
+sync_batchnorm: True
+accelerator: "gpu"
+strategy: "ddp"
+precision: 16-mixed

solo/losses/__init__.py

@@ -32,6 +32,7 @@
 from solo.losses.vibcreg import vibcreg_loss_func
 from solo.losses.vicreg import vicreg_loss_func
 from solo.losses.wmse import wmse_loss_func
+from solo.losses.ssley import ssley_loss_func
 
 __all__ = [
     "barlow_loss_func",
@@ -49,4 +50,5 @@
     "vibcreg_loss_func",
     "vicreg_loss_func",
     "wmse_loss_func",
+    "ssley_loss_func"
 ]

solo/losses/ssley.py

@@ -0,0 +1,53 @@
+# Copyright 2023 solo-learn development team.
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to use,
+# copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+# Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all copies
+# or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
+# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
+# FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+import torch
+import torch.nn.functional as F
+from solo.utils.misc import gather
+
+
+def ssley_loss_func(
+    z1: torch.Tensor,
+    z2: torch.Tensor,
+) -> torch.Tensor:
+    """Computes SSL-EY's loss given batch of projected features z1 from view 1 and
+    projected features z2 from view 2.
+
+    Args:
+        z1 (torch.Tensor): NxD Tensor containing projected features from view 1.
+        z2 (torch.Tensor): NxD Tensor containing projected features from view 2.
+
+    Returns:
+        torch.Tensor: VICReg loss.
+    """
+
+    sim_loss = invariance_loss(z1, z2)
+
+    N, D = z1.size()
+    B = torch.cov(torch.hstack((z1, z2)).T)
+
+    if dist.is_available() and dist.is_initialized():
+        dist.all_reduce(B)
+        world_size = dist.get_world_size()
+        B /= world_size
+
+    A = B[:D, D:] + B[D:, :D]
+    B = B[:D, :D] + B[D:, D:]
+
+    return -torch.trace(2 * A - B @ B)

solo/methods/__init__.py

@@ -37,6 +37,7 @@
 from solo.methods.vibcreg import VIbCReg
 from solo.methods.vicreg import VICReg
 from solo.methods.wmse import WMSE
+from solo.methods.ssley import SSLEY
 
 METHODS = {
     # base classes
@@ -61,6 +62,7 @@
     "vibcreg": VIbCReg,
     "vicreg": VICReg,
     "wmse": WMSE,
+    "ssley": SSLEY,
 }
 __all__ = [
     "BarlowTwins",
@@ -83,4 +85,5 @@
     "VIbCReg",
     "VICReg",
     "WMSE",
+    "SSLEY",
 ]