added swin swav model

pl_bolts/models/self_supervised/swav/__init__.py

@@ -1,6 +1,7 @@
 from pl_bolts.models.self_supervised.swav.loss import SWAVLoss
 from pl_bolts.models.self_supervised.swav.swav_module import SwAV
 from pl_bolts.models.self_supervised.swav.swav_resnet import resnet18, resnet50
+from pl_bolts.models.self_supervised.swav.swav_swin import swin_s , swin_b , swin_v2_t , swin_v2_s , swin_v2_b
 from pl_bolts.models.self_supervised.swav.transforms import (
     SwAVEvalDataTransform,
     SwAVFinetuneTransform,
@@ -11,6 +12,11 @@
     "SwAV",
     "resnet18",
     "resnet50",
+    "swin_s",
+    "swin_b",
+    "swin_v2_t",
+    "swin_v2_s",
+    "swin_v2_b",
     "SwAVEvalDataTransform",
     "SwAVFinetuneTransform",
     "SwAVTrainDataTransform",

pl_bolts/models/self_supervised/swav/swav_module.py

@@ -9,6 +9,8 @@
 
 from pl_bolts.models.self_supervised.swav.loss import SWAVLoss
 from pl_bolts.models.self_supervised.swav.swav_resnet import resnet18, resnet50
+from pl_bolts.models.self_supervised.swav.swav_swin import swin_s , swin_b , swin_v2_t , swin_v2_s , swin_v2_b
+
 from pl_bolts.optimizers.lars import LARS
 from pl_bolts.optimizers.lr_scheduler import linear_warmup_decay
 from pl_bolts.transforms.dataset_normalizations import (
@@ -158,6 +160,16 @@ def init_model(self):
             backbone = resnet18
         elif self.arch == "resnet50":
             backbone = resnet50
+        elif self.arch == "swin_s":
+            backbone = swin_s
+        elif self.arch == "swin_b":
+            backbone = swin_b
+        elif self.arch == "swin_v2_t":
+            backbone = swin_v2_t
+        elif self.arch == "swin_v2_s":
+            backbone = swin_v2_s
+        elif self.arch == "swin_v2_b":
+            backbone = swin_v2_b
 
         return backbone(
             normalize=True,