-
Notifications
You must be signed in to change notification settings - Fork 342
[WIP] Implement SSL-EY #1443
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
jameschapman19
wants to merge
26
commits into
lightly-ai:master
Choose a base branch
from
jameschapman19:master
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
[WIP] Implement SSL-EY #1443
Changes from 7 commits
Commits
Show all changes
26 commits
Select commit
Hold shift + click to select a range
1ce3265
Adding SSL-EY
jameschapman19 c8943b1
Adding SSL-EY
jameschapman19 410d971
Adding SSL-EY to docs
jameschapman19 fbd6438
Adding SSL-EY to docs
jameschapman19 d41c0f7
Adding SSL-EY to tests
jameschapman19 6dbb301
Adding SSL-EY to tests
jameschapman19 ea0a3a0
Adding SSL-EY to tests
jameschapman19 3aae9fb
Adding SSL-EY to tests
jameschapman19 7ebd395
Adding SSL-EY to tests
jameschapman19 0a08b0b
Adding SSL-EY to tests
jameschapman19 e301388
Merge branch 'master' into add-ssl-ey
551d71b
Format
bf3bab6
Cleanup types
1c86d1b
Cleanup tests
6f82ba3
Use SSLEYTransform
b8f4920
Fix types
66d2787
Add SSLEYTransform test
62cf5b4
Add SSLEYProjectionHead
6b9bc35
Update examples
b4bfa87
Fix citation
1f4c5a8
Fix typings in dist.py
7e3720f
Format
b9d0ed9
Remove unused import
9261583
Use any instead of FunctionCtx
e8e3054
Add type ignore override
f88efd0
Show SSL-EY in docs
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 <https://arxiv.org/abs/2103.03230>`_ 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 <https://arxiv.org/abs/2310.01012>`_ | ||
|
|
||
|
|
||
| .. 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 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This import doesn't work as
ssley_transformdoesn't exist yet. I think it would make sense to add assley_transform.pymodule and create aSSLEYTransformwhich inherits fromVICRegTransform. That way one doesn't have to remember that ssley uses the vicreg transform.