Repository: lightly-ai/lightly

Issue:
Refactor CLI (deprecate high-level models)
# Refactor CLI (deprecate high-level models)

With `1.1.2` we will deprecate all mid- and high-level models (see #498). This issue is about removing the high-level model from the CLI and implementing an alternative.

- [ ] Copy `SimCLR` into the cli part as an internal module and use it from there
- [ ] Remove the dependence of the train and embed cli on the `models/simclr.py`
- [ ] Make sure old checkpoint still work 


Discussion:

Resolving patch:
diff --git a/lightly/cli/_cli_simclr.py b/lightly/cli/_cli_simclr.py
new file mode 100644
--- /dev/null
+++ b/lightly/cli/_cli_simclr.py
@@ -0,0 +1,49 @@
+""" SimCLR Model """
+
+# Copyright (c) 2020. Lightly AG and its affiliates.
+# All Rights Reserved
+
+import warnings
+
+import torch
+import torch.nn as nn
+
+from lightly.models.modules import SimCLRProjectionHead
+
+
+class _SimCLR(nn.Module):
+    """Implementation of SimCLR used by the command-line interface.
+
+        Provides backwards compatability with old checkpoints.
+    """
+
+    def __init__(self, backbone: nn.Module, num_ftrs: int = 32,
+                 out_dim: int = 128):
+
+        super(_SimCLR, self).__init__()
+
+        self.backbone = backbone
+        self.projection_head = SimCLRProjectionHead(num_ftrs, num_ftrs,
+                                                    out_dim)
+
+
+    def forward(self, x0: torch.Tensor, x1: torch.Tensor = None):
+        """Embeds and projects the input images.
+
+        """
+
+        # forward pass of first input x0
+        f0 = self.backbone(x0).flatten(start_dim=1)
+        out0 = self.projection_head(f0)
+
+
+        # return out0 if x1 is None
+        if x1 is None:
+            return out0
+
+        # forward pass of second input x1
+        f1 = self.backbone(x1).flatten(start_dim=1)
+        out1 = self.projection_head(f1)
+
+        # return both outputs
+        return out0, out1
\ No newline at end of file
diff --git a/lightly/cli/embed_cli.py b/lightly/cli/embed_cli.py
--- a/lightly/cli/embed_cli.py
+++ b/lightly/cli/embed_cli.py
@@ -16,9 +16,9 @@
 import torchvision
 from torch.utils.hipify.hipify_python import bcolors
 
+from lightly.cli._cli_simclr import _SimCLR
 from lightly.data import LightlyDataset
 from lightly.embedding import SelfSupervisedEmbedding
-from lightly.models import SimCLR
 
 from lightly.models import ResNetGenerator
 from lightly.models.batchnorm import get_norm_layer
@@ -100,7 +100,7 @@ def _embed_cli(cfg, is_cli_call=True):
         nn.AdaptiveAvgPool2d(1),
     )
 
-    model = SimCLR(
+    model = _SimCLR(
         features,
         num_ftrs=cfg['model']['num_ftrs'],
         out_dim=cfg['model']['out_dim']
diff --git a/lightly/cli/train_cli.py b/lightly/cli/train_cli.py
--- a/lightly/cli/train_cli.py
+++ b/lightly/cli/train_cli.py
@@ -15,11 +15,11 @@
 
 from torch.utils.hipify.hipify_python import bcolors
 
+from lightly.cli._cli_simclr import _SimCLR
 from lightly.data import ImageCollateFunction
 from lightly.data import LightlyDataset
 from lightly.embedding import SelfSupervisedEmbedding
 from lightly.loss import NTXentLoss
-from lightly.models import SimCLR
 
 from lightly.models import ResNetGenerator
 from lightly.models.batchnorm import get_norm_layer
@@ -103,7 +103,7 @@ def _train_cli(cfg, is_cli_call=True):
         nn.AdaptiveAvgPool2d(1),
     )
 
-    model = SimCLR(
+    model = _SimCLR(
         features,
         num_ftrs=cfg['model']['num_ftrs'],
         out_dim=cfg['model']['out_dim']


Test patch: