simclr fixes

docs/source/self_supervised_models.rst

@@ -274,6 +274,7 @@ Model implemented by:
 
     - `William Falcon <https://github.com/williamFalcon>`_
     - `Tullie Murrell <https://github.com/tullie>`_
+    - `Ananya Harsh Jha <https://github.com/ananyahjha93>`_
 
 To Train::
 
@@ -289,7 +290,7 @@ To Train::
     dm.val_transforms = SimCLREvalDataTransform(32)
 
     # model
-    model = SimCLR(num_samples=dm.num_samples, batch_size=dm.batch_size)
+    model = SimCLR(num_samples=dm.num_samples, batch_size=dm.batch_size, dataset='cifar10')
 
     # fit
     trainer = pl.Trainer()
@@ -310,21 +311,29 @@ CIFAR-10 baseline
      - Hardware
      - LR
    * - `Original <https://github.com/google-research/simclr#finetuning-the-linear-head-linear-eval>`_
-     - `92.00? <https://github.com/google-research/simclr#finetuning-the-linear-head-linear-eval>`_
+     - `~94.00 <https://github.com/google-research/simclr#finetuning-the-linear-head-linear-eval>`_
      - resnet50
      - LARS
-     - 512
-     - 1000
-     - 1 V100 (32GB)
-     - 1.0
+     - 2048
+     - 800
+     - TPUs
+     - 1.0/1.5
    * - Ours
      - `85.68 <https://tensorboard.dev/experiment/GlS1eLXMQsqh3T5DAec6UQ/#scalars>`_
      - `resnet50 <https://github.com/PyTorchLightning/PyTorch-Lightning-Bolts/blob/master/pl_bolts/models/self_supervised/resnets.py#L301-L309>`_
-     - `LARS <https://pytorch-lightning-bolts.readthedocs.io/en/latest/api/pl_bolts.optimizers.lars_scheduling.html#pl_bolts.optimizers.lars_scheduling.LARSWrapper>`_
-     - 512
-     - 960 (12 hr)
-     - 1 V100 (32GB)
-     - 1e-6
+     - `LARS-SGD <https://pytorch-lightning-bolts.readthedocs.io/en/latest/api/pl_bolts.optimizers.lars_scheduling.html#pl_bolts.optimizers.lars_scheduling.LARSWrapper>`_
+     - 2048
+     - 800 (~4 hours)
+     - 8 V100 (16GB)
+     - 1.5
+   * - Ours
+     - `85.68 <https://tensorboard.dev/experiment/GlS1eLXMQsqh3T5DAec6UQ/#scalars>`_
+     - `resnet50 <https://github.com/PyTorchLightning/PyTorch-Lightning-Bolts/blob/master/pl_bolts/models/self_supervised/resnets.py#L301-L309>`_
+     - `LARS-Adam <https://pytorch-lightning-bolts.readthedocs.io/en/latest/api/pl_bolts.optimizers.lars_scheduling.html#pl_bolts.optimizers.lars_scheduling.LARSWrapper>`_
+     - 2048
+     - 800 (~4 hours)
+     - 8 V100 (16GB)
+     - 1e-3
 
 |
 
@@ -422,7 +431,7 @@ To Train::
     model = SwAV(
         gpus=1,
         num_samples=dm.num_unlabeled_samples,
-        datamodule=dm,
+        dataset='stl10,
         batch_size=batch_size
     )
 

pl_bolts/callbacks/ssl_online.py

@@ -26,14 +26,15 @@ class SSLOnlineEvaluator(Callback):  # pragma: no-cover
     """
     def __init__(
         self,
+        dataset: str,
         drop_p: float = 0.2,
         hidden_dim: Optional[int] = None,
         z_dim: int = None,
         num_classes: int = None,
-        dataset: str = 'stl10'
     ):
         """
         Args:
+            dataset: if stl10, need to get the labeled batch
             drop_p: Dropout probability
             hidden_dim: Hidden dimension for the fine-tune MLP
             z_dim: Representation dimension

pl_bolts/models/self_supervised/byol/byol_module.py

@@ -105,7 +105,7 @@ def cosine_similarity(self, a, b):
         return sim
 
     def shared_step(self, batch, batch_idx):
-        (img_1, img_2), y = batch
+        (img_1, img_2, _), y = batch
 
         # Image 1 to image 2 loss
         y1, z1, h1 = self.online_network(img_1)
@@ -220,15 +220,8 @@ def cli_main():
 
     model = BYOL(**args.__dict__)
 
-    def to_device(batch, device):
-        (x1, x2), y = batch
-        x1 = x1.to(device)
-        y = y.to(device)
-        return x1, y
-
     # finetune in real-time
-    online_eval = SSLOnlineEvaluator(z_dim=2048, num_classes=dm.num_classes)
-    online_eval.to_device = to_device
+    online_eval = SSLOnlineEvaluator(dataset=args.dataset, z_dim=2048, num_classes=dm.num_classes)
 
     trainer = pl.Trainer.from_argparse_args(args, max_steps=300000, callbacks=[online_eval])
 

pl_bolts/models/self_supervised/byol/models.py

@@ -27,17 +27,13 @@ def __init__(self, encoder=None):
             encoder = torchvision_ssl_encoder('resnet50')
         # Encoder
         self.encoder = encoder
-        # Pooler
-        self.pooler = nn.AdaptiveAvgPool2d((1, 1))
         # Projector
         self.projector = MLP()
         # Predictor
         self.predictor = MLP(input_dim=256)
 
     def forward(self, x):
         y = self.encoder(x)[0]
-        y = self.pooler(y)
-        y = y.view(y.size(0), -1)
         z = self.projector(y)
         h = self.predictor(z)
         return y, z, h

pl_bolts/models/self_supervised/resnets.py

@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn
 
 from pl_bolts.utils.warnings import warn_missing_pkg
@@ -12,7 +13,6 @@
     'resnet18',
     'resnet34',
     'resnet50',
-    'resnet50_bn',
     'resnet101',
     'resnet152',
     'resnext50_32x4d',
@@ -25,7 +25,7 @@
 MODEL_URLS = {
     'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
     'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
-    'resnet50_bn': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
     'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
     'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
     'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
@@ -92,49 +92,6 @@ class Bottleneck(nn.Module):
     def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                  base_width=64, dilation=1, norm_layer=None):
         super(Bottleneck, self).__init__()
-        # if norm_layer is None:
-        #     norm_layer = nn.BatchNorm2d
-        width = int(planes * (base_width / 64.)) * groups
-        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
-        self.conv1 = conv1x1(inplanes, width)
-        # self.bn1 = norm_layer(width)
-        self.conv2 = conv3x3(width, width, stride, groups, dilation)
-        # self.bn2 = norm_layer(width)
-        self.conv3 = conv1x1(width, planes * self.expansion)
-        # self.bn3 = norm_layer(planes * self.expansion)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        identity = x
-
-        out = self.conv1(x)
-        # out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        # out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        # out = self.bn3(out)
-
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        out += identity
-        out = self.relu(out)
-
-        return out
-
-
-class BottleneckBN(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
-                 base_width=64, dilation=1, norm_layer=None):
-        super(BottleneckBN, self).__init__()
         if norm_layer is None:
             norm_layer = nn.BatchNorm2d
         width = int(planes * (base_width / 64.)) * groups
@@ -174,9 +131,20 @@ def forward(self, x):
 
 class ResNet(nn.Module):
 
-    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
-                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
-                 norm_layer=None, return_all_feature_maps=False):
+    def __init__(
+        self,
+        block,
+        layers,
+        num_classes=1000,
+        zero_init_residual=False,
+        groups=1,
+        width_per_group=64,
+        replace_stride_with_dilation=None,
+        norm_layer=None,
+        return_all_feature_maps=False,
+        first_conv=True,
+        maxpool1=True,
+    ):
         super(ResNet, self).__init__()
         if norm_layer is None:
             norm_layer = nn.BatchNorm2d
@@ -194,11 +162,24 @@ def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
                              "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
         self.groups = groups
         self.base_width = width_per_group
-        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
-                               bias=False)
+
+        if first_conv:
+            self.conv1 = nn.Conv2d(
+                3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
+            )
+        else:
+            self.conv1 = nn.Conv2d(
+                3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False
+            )
+
         self.bn1 = norm_layer(self.inplanes)
         self.relu = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        if maxpool1:
+            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        else:
+            self.maxpool = nn.MaxPool2d(kernel_size=1, stride=1)
+
         self.layer1 = self._make_layer(block, 64, layers[0])
         self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
                                        dilate=replace_stride_with_dilation[0])
@@ -269,6 +250,9 @@ def forward(self, x):
             x0 = self.layer3(x0)
             x0 = self.layer4(x0)
 
+            x0 = self.avgpool(x0)
+            x0 = torch.flatten(x0, 1)
+
             return [x0]
 
 
@@ -314,17 +298,6 @@ def resnet50(pretrained: bool = False, progress: bool = True, **kwargs):
     return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
 
 
-def resnet50_bn(pretrained: bool = False, progress: bool = True, **kwargs):
-    r"""ResNet-50 model from
-    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
-
-    Args:
-        pretrained: If True, returns a model pre-trained on ImageNet
-        progress: If True, displays a progress bar of the download to stderr
-    """
-    return _resnet('resnet50_bn', BottleneckBN, [3, 4, 6, 3], pretrained, progress, **kwargs)
-
-
 def resnet101(pretrained: bool = False, progress: bool = True, **kwargs):
     r"""ResNet-101 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
@@ -333,7 +306,7 @@ def resnet101(pretrained: bool = False, progress: bool = True, **kwargs):
         pretrained: If True, returns a model pre-trained on ImageNet
         progress: If True, displays a progress bar of the download to stderr
     """
-    return _resnet('resnet101', BottleneckBN, [3, 4, 23, 3], pretrained, progress, **kwargs)
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
 
 
 def resnet152(pretrained: bool = False, progress: bool = True, **kwargs):