Revision of the MoCo SSL model (#928)

Lightning-Universe · Jul 10, 2023 · a4a64af · a4a64af
1 parent 899f40d
commit a4a64af
Show file tree

Hide file tree

Showing 13 changed files with 506 additions and 417 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,7 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
--
+- Revision of the MoCo SSL model ([#928](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/928))
 
 
 ### Deprecated

diff --git a/docs/source/introduction_guide.rst b/docs/source/introduction_guide.rst
@@ -22,7 +22,7 @@ All models are tested (daily), benchmarked, documented and work on CPUs, TPUs, G
 
     from pl_bolts.models import VAE
     from pl_bolts.models.vision import GPT2, ImageGPT, PixelCNN
-    from pl_bolts.models.self_supervised import AMDIM, CPC_v2, SimCLR, Moco_v2
+    from pl_bolts.models.self_supervised import AMDIM, CPC_v2, SimCLR, MoCo
     from pl_bolts.models import LinearRegression, LogisticRegression
     from pl_bolts.models.gans import GAN
     from pl_bolts.callbacks import PrintTableMetricsCallback

diff --git a/docs/source/models/self_supervised.rst b/docs/source/models/self_supervised.rst
@@ -259,7 +259,7 @@ CPC (v2) API
 Moco (v2) API
 ^^^^^^^^^^^^^
 
-.. autoclass:: pl_bolts.models.self_supervised.Moco_v2
+.. autoclass:: pl_bolts.models.self_supervised.MoCo
    :noindex:
 
 SimCLR

diff --git a/pyproject.toml b/pyproject.toml
@@ -198,7 +198,6 @@ module = [
     "pl_bolts.models.self_supervised.cpc.transforms",
     "pl_bolts.models.self_supervised.evaluator",
     "pl_bolts.models.self_supervised.moco.callbacks",
-    "pl_bolts.models.self_supervised.moco.moco2_module",
     "pl_bolts.models.self_supervised.moco.transforms",
     "pl_bolts.models.self_supervised.resnets",
     "pl_bolts.models.self_supervised.simclr.simclr_finetuner",

diff --git a/src/pl_bolts/models/detection/retinanet/retinanet_module.py b/src/pl_bolts/models/detection/retinanet/retinanet_module.py
@@ -135,10 +135,7 @@ def configure_optimizers(self):
 
 @under_review()
 def cli_main():
-    try:  # Backward compatibility for Lightning CLI
-        from pytorch_lightning.cli import LightningCLI  # PL v1.9+
-    except ImportError:
-        from pytorch_lightning.utilities.cli import LightningCLI  # PL v1.8
+    from pytorch_lightning.cli import LightningCLI
 
     from pl_bolts.datamodules import VOCDetectionDataModule
 

diff --git a/src/pl_bolts/models/detection/yolo/yolo_module.py b/src/pl_bolts/models/detection/yolo/yolo_module.py
@@ -4,7 +4,6 @@
 import torch
 import torch.nn as nn
 from pytorch_lightning import LightningModule
-from pytorch_lightning.utilities.cli import LightningCLI
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 from torch import Tensor, optim
 
@@ -614,4 +613,6 @@ def _resize(self, image: Tensor, target: TARGET) -> Tuple[Tensor, TARGET]:
 
 
 if __name__ == "__main__":
+    from pytorch_lightning.cli import LightningCLI
+
     LightningCLI(CLIYOLO, ResizedVOCDetectionDataModule, seed_everything_default=42)
diff --git a/src/pl_bolts/models/self_supervised/__init__.py b/src/pl_bolts/models/self_supervised/__init__.py
@@ -20,7 +20,7 @@
 from pl_bolts.models.self_supervised.byol.byol_module import BYOL
 from pl_bolts.models.self_supervised.cpc.cpc_module import CPC_v2
 from pl_bolts.models.self_supervised.evaluator import SSLEvaluator
-from pl_bolts.models.self_supervised.moco.moco2_module import Moco_v2
+from pl_bolts.models.self_supervised.moco.moco_module import MoCo
 from pl_bolts.models.self_supervised.simclr.simclr_module import SimCLR
 from pl_bolts.models.self_supervised.simsiam.simsiam_module import SimSiam
 from pl_bolts.models.self_supervised.ssl_finetuner import SSLFineTuner
@@ -31,7 +31,7 @@
     "BYOL",
     "CPC_v2",
     "SSLEvaluator",
-    "Moco_v2",
+    "MoCo",
     "SimCLR",
     "SimSiam",
     "SSLFineTuner",

diff --git a/src/pl_bolts/models/self_supervised/moco/callbacks.py b/src/pl_bolts/models/self_supervised/moco/callbacks.py
@@ -6,7 +6,7 @@
 
 
 @under_review()
-class MocoLRScheduler(Callback):
+class MoCoLRScheduler(Callback):
     def __init__(self, initial_lr=0.03, use_cosine_scheduler=False, schedule=(120, 160), max_epochs=200) -> None:
         super().__init__()
         self.lr = initial_lr