Refactor the rest of pl_bolts.models.self_supervised

pl_bolts/models/self_supervised/__init__.py

@@ -26,3 +26,15 @@
 from pl_bolts.models.self_supervised.simclr.simclr_module import SimCLR  # noqa: F401
 from pl_bolts.models.self_supervised.ssl_finetuner import SSLFineTuner  # noqa: F401
 from pl_bolts.models.self_supervised.swav.swav_module import SwAV  # noqa: F401
+
+
+__all__ = [
+    "AMDIM",
+    "BYOL",
+    "CPCV2",
+    "SSLEvaluator"
+    "MocoV2",
+    "SimCLR",
+    "SSLFineTuner",
+    "SwAV",
+]

pl_bolts/models/self_supervised/cpc/__init__.py

@@ -8,3 +8,16 @@
     CPCTrainTransformsImageNet128,
     CPCTrainTransformsSTL10,
 )
+
+
+__all__ = [
+    "CPCV2",
+    "cpc_resnet50",
+    "cpc_resnet101",
+    "CPCEvalTransformsCIFAR10",
+    "CPCEvalTransformsImageNet128",
+    "CPCEvalTransformsSTL10",
+    "CPCTrainTransformsCIFAR10",
+    "CPCTrainTransformsImageNet128",
+    "CPCTrainTransformsSTL10",
+]

pl_bolts/models/self_supervised/cpc/transforms.py

@@ -4,8 +4,8 @@
 
 if _TORCHVISION_AVAILABLE:
     from torchvision import transforms
-else:
-    warn_missing_pkg('torchvision')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg('torchvision')
 
 
 class CPCTrainTransformsCIFAR10:
@@ -39,8 +39,8 @@ def __init__(self, patch_size=8, overlap=4):
             patch_size: size of patches when cutting up the image into overlapping patches
             overlap: how much to overlap patches
         """
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -97,8 +97,8 @@ def __init__(self, patch_size: int = 8, overlap: int = 4):
             patch_size: size of patches when cutting up the image into overlapping patches
             overlap: how much to overlap patches
         """
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -151,8 +151,8 @@ def __init__(self, patch_size: int = 16, overlap: int = 8):
             patch_size: size of patches when cutting up the image into overlapping patches
             overlap: how much to overlap patches
         """
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -210,8 +210,8 @@ def __init__(self, patch_size: int = 16, overlap: int = 8):
             patch_size: size of patches when cutting up the image into overlapping patches
             overlap: how much to overlap patches
         """
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -260,8 +260,8 @@ def __init__(self, patch_size: int = 32, overlap: int = 16):
             patch_size: size of patches when cutting up the image into overlapping patches
             overlap: how much to overlap patches
         """
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -320,8 +320,8 @@ def __init__(self, patch_size: int = 32, overlap: int = 16):
             patch_size: size of patches when cutting up the image into overlapping patches
             overlap: how much to overlap patches
         """
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 

pl_bolts/models/self_supervised/moco/transforms.py

@@ -10,13 +10,13 @@
 
 if _TORCHVISION_AVAILABLE:
     from torchvision import transforms
-else:
-    warn_missing_pkg('torchvision')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg('torchvision')
 
 if _PIL_AVAILABLE:
     from PIL import ImageFilter
-else:
-    warn_missing_pkg('PIL', pypi_name='Pillow')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg('PIL', pypi_name='Pillow')
 
 
 class Moco2TrainCIFAR10Transforms:
@@ -25,8 +25,8 @@ class Moco2TrainCIFAR10Transforms:
     https://arxiv.org/pdf/2003.04297.pdf
     """
     def __init__(self, height: int = 32):
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -55,8 +55,8 @@ class Moco2EvalCIFAR10Transforms:
     https://arxiv.org/pdf/2003.04297.pdf
     """
     def __init__(self, height: int = 32):
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -79,8 +79,8 @@ class Moco2TrainSTL10Transforms:
     https://arxiv.org/pdf/2003.04297.pdf
     """
     def __init__(self, height: int = 64):
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -109,8 +109,8 @@ class Moco2EvalSTL10Transforms:
     https://arxiv.org/pdf/2003.04297.pdf
     """
     def __init__(self, height: int = 64):
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -135,8 +135,8 @@ class Moco2TrainImagenetTransforms:
     """
 
     def __init__(self, height: int = 128):
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -165,8 +165,8 @@ class Moco2EvalImagenetTransforms:
     https://arxiv.org/pdf/2003.04297.pdf
     """
     def __init__(self, height: int = 128):
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -187,8 +187,8 @@ class GaussianBlur(object):
     """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709"""
 
     def __init__(self, sigma=(0.1, 2.0)):
-        if not _PIL_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _PIL_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `Pillow` which is not installed yet, install it with `pip install Pillow`.'
             )
         self.sigma = sigma

pl_bolts/models/self_supervised/simclr/simclr_finetuner.py

@@ -13,7 +13,7 @@
 )
 
 
-def cli_main():  # pragma: no-cover
+def cli_main():  # pragma: no cover
     from pl_bolts.datamodules import CIFAR10DataModule, ImagenetDataModule, STL10DataModule
 
     pl.seed_everything(1234)

pl_bolts/models/self_supervised/simclr/transforms.py

@@ -1,17 +1,17 @@
 import numpy as np
 
-from pl_bolts.utils import _TORCHVISION_AVAILABLE
+from pl_bolts.utils import _TORCHVISION_AVAILABLE, _OPENCV_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
 if _TORCHVISION_AVAILABLE:
     import torchvision.transforms as transforms
 else:
     warn_missing_pkg('torchvision')  # pragma: no-cover
 
-try:
+if _OPENCV_AVAILABLE:
     import cv2
-except ModuleNotFoundError:
-    warn_missing_pkg('cv2', pypi_name='opencv-python')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg('cv2', pypi_name='opencv-python')
 
 
 class SimCLRTrainDataTransform(object):
@@ -43,8 +43,8 @@ def __init__(
         normalize=None
     ) -> None:
 
-        if not _TORCHVISION_AVAILABLE:
-            raise ModuleNotFoundError(  # pragma: no-cover
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
                 'You want to use `transforms` from `torchvision` which is not installed yet.'
             )
 
@@ -187,6 +187,11 @@ def __call__(self, sample):
 class GaussianBlur(object):
     # Implements Gaussian blur as described in the SimCLR paper
     def __init__(self, kernel_size, p=0.5, min=0.1, max=2.0):
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
+                'You want to use `GaussianBlur` from `cv2` which is not installed yet.'
+            )
+
         self.min = min
         self.max = max
 

pl_bolts/models/self_supervised/swav/__init__.py

@@ -5,3 +5,13 @@
     SwAVFinetuneTransform,
     SwAVTrainDataTransform,
 )
+
+
+__all__ = [
+    "SwAV",
+    "resnet18",
+    "resnet50",
+    "SwAVEvalDataTransform",
+    "SwAVFinetuneTransform",
+    "SwAVTrainDataTransform",
+]

pl_bolts/models/self_supervised/swav/swav_finetuner.py

@@ -9,7 +9,7 @@
 from pl_bolts.transforms.dataset_normalizations import imagenet_normalization, stl10_normalization
 
 
-def cli_main():  # pragma: no-cover
+def cli_main():  # pragma: no cover
     from pl_bolts.datamodules import ImagenetDataModule, STL10DataModule
 
     pl.seed_everything(1234)

pl_bolts/models/self_supervised/swav/transforms.py

@@ -7,13 +7,13 @@
 
 if _TORCHVISION_AVAILABLE:
     import torchvision.transforms as transforms
-else:
-    warn_missing_pkg('torchvision')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg('torchvision')
 
 if _OPENCV_AVAILABLE:
     import cv2
-else:
-    warn_missing_pkg('cv2', pypi_name='opencv-python')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg('cv2', pypi_name='opencv-python')
 
 
 class SwAVTrainDataTransform(object):