Moco

passl/data/preprocess/basic_transforms.py

@@ -57,6 +57,7 @@
     "SimCLRGaussianBlur",
     "BYOLSolarize",
     "MAERandCropImage",
+    "GaussianBlur",
 ]
 
 
@@ -941,3 +942,21 @@ def __call__(self, img):
             else:
                 img = ImageOps.solarize(img)
         return img
+
+class GaussianBlur(object):
+    """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709"""
+
+    def __init__(self, sigma=[.1, 2.], p=1.0):
+        self.p = p
+        self.sigma = sigma
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            if not isinstance(img, Image.Image):
+                img = np.ascontiguousarray(img)
+                img = Image.fromarray(img)
+            sigma = random.uniform(self.sigma[0], self.sigma[1])
+            img = img.filter(ImageFilter.GaussianBlur(radius=sigma))
+            if isinstance(img, Image.Image):
+                img = np.asarray(img)
+        return img

passl/engine/loops/contrastive_learning_loop.py

@@ -68,6 +68,9 @@ def train_one_step(self, batch):
         # remove label
         batch = batch[0]
 
+        for idx, value in enumerate(batch):
+            if isinstance(value,paddle.Tensor):
+                batch[idx] = batch[idx].cuda()
         # do forward and backward
         loss_dict = self.forward_backward(batch)
 

passl/engine/loops/loop.py

@@ -286,6 +286,7 @@ def train_one_epoch(self):
                     paddle.to_tensor(batch[0]['label'])
                 ]
 
+
             self.global_step += 1
 
             # do forward and backward

passl/models/__init__.py

@@ -29,7 +29,7 @@
 from .mocov3 import *
 from .swav import *
 from .simsiam import *
-
+from .mocov2 import *
 __all__ = ["build_model"]