Revision of the MoCo SSL model

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

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

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",

src/pl_bolts/models/detection/retinanet/retinanet_module.py

@@ -136,13 +136,17 @@ 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
+        import pytorch_lightning.cli
+
+        cli_class: Any = getattr(pytorch_lightning.cli, "LightningCLI")  # PL v1.9+
+    except Exception:
+        import pytorch_lightning.utilities.cli
+
+        cli_class = getattr(pytorch_lightning.utilities.cli, "LightningCLI")  # PL v1.8
 
     from pl_bolts.datamodules import VOCDetectionDataModule
 
-    LightningCLI(RetinaNet, VOCDetectionDataModule, seed_everything_default=42)
+    cli_class(RetinaNet, VOCDetectionDataModule, seed_everything_default=42)
 
 
 if __name__ == "__main__":

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,13 @@ def _resize(self, image: Tensor, target: TARGET) -> Tuple[Tensor, TARGET]:
 
 
 if __name__ == "__main__":
-    LightningCLI(CLIYOLO, ResizedVOCDetectionDataModule, seed_everything_default=42)
+    try:  # Backward compatibility for Lightning CLI
+        import pytorch_lightning.cli
+
+        cli_class: Any = getattr(pytorch_lightning.cli, "LightningCLI")  # PL v1.9+
+    except Exception:
+        import pytorch_lightning.utilities.cli
+
+        cli_class = getattr(pytorch_lightning.utilities.cli, "LightningCLI")  # PL v1.8
+
+    cli_class(CLIYOLO, ResizedVOCDetectionDataModule, seed_everything_default=42)

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",

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