Refactor pl_bolts.models.self_supervised.amdim

pl_bolts/models/self_supervised/amdim/__init__.py

@@ -1,14 +1,21 @@
 from pl_bolts.models.self_supervised.amdim.amdim_module import AMDIM  # noqa: F401
 from pl_bolts.models.self_supervised.amdim.networks import AMDIMEncoder  # noqa: F401
+from pl_bolts.models.self_supervised.amdim.transforms import (  # noqa: F401
+    AMDIMEvalTransformsCIFAR10,
+    AMDIMEvalTransformsImageNet128,
+    AMDIMEvalTransformsSTL10,
+    AMDIMTrainTransformsCIFAR10,
+    AMDIMTrainTransformsImageNet128,
+    AMDIMTrainTransformsSTL10,
+)
 
-try:
-    from pl_bolts.models.self_supervised.amdim.transforms import (  # noqa: F401
-        AMDIMEvalTransformsCIFAR10,
-        AMDIMEvalTransformsImageNet128,
-        AMDIMEvalTransformsSTL10,
-        AMDIMTrainTransformsCIFAR10,
-        AMDIMTrainTransformsImageNet128,
-        AMDIMTrainTransformsSTL10,
-    )
-except ModuleNotFoundError:
-    pass
+__all__ = [
+    "AMDIM",
+    "AMDIMEncoder",
+    "AMDIMEvalTransformsCIFAR10",
+    "AMDIMEvalTransformsImageNet128",
+    "AMDIMEvalTransformsSTL10",
+    "AMDIMTrainTransformsCIFAR10",
+    "AMDIMTrainTransformsImageNet128",
+    "AMDIMTrainTransformsSTL10",
+]

pl_bolts/models/self_supervised/amdim/amdim_module.py

@@ -7,11 +7,8 @@
 import torch.optim as optim
 from torch.utils.data import DataLoader
 
-try:
-    from pl_bolts.models.self_supervised.amdim.datasets import AMDIMPretraining
-except ModuleNotFoundError:
-    pass
 from pl_bolts.losses.self_supervised_learning import FeatureMapContrastiveTask
+from pl_bolts.models.self_supervised.amdim.datasets import AMDIMPretraining
 from pl_bolts.models.self_supervised.amdim.networks import AMDIMEncoder
 from pl_bolts.utils.self_supervised import torchvision_ssl_encoder
 

pl_bolts/models/self_supervised/amdim/datasets.py

@@ -2,16 +2,16 @@
 
 from torch.utils.data import random_split
 
+from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet
+from pl_bolts.datasets.ssl_amdim_datasets import CIFAR10Mixed
+from pl_bolts.models.self_supervised.amdim import transforms as amdim_transforms
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-try:
+if _TORCHVISION_AVAILABLE:
     from torchvision.datasets import STL10
-
-    from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet
-    from pl_bolts.datasets.ssl_amdim_datasets import CIFAR10Mixed
-    from pl_bolts.models.self_supervised.amdim import transforms as amdim_transforms
-except ModuleNotFoundError:
-    warn_missing_pkg('torchvision')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg('torchvision')
 
 
 class AMDIMPretraining():
@@ -55,6 +55,10 @@ def imagenet(dataset_root, nb_classes, split: str = 'train'):
 
     @staticmethod
     def stl(dataset_root, split: Optional[str] = None):
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
+                'You want to use STL10 dataset loaded from `torchvision` which is not installed yet.'
+            )
         dataset = STL10(
             root=dataset_root,
             split='unlabeled',
@@ -96,6 +100,10 @@ def cifar10(dataset_root, patch_size, patch_overlap, split: str = 'train'):
 
     @staticmethod
     def stl(dataset_root, patch_size, patch_overlap, split: Optional[str] = None):
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
+                'You want to use STL10 dataset loaded from `torchvision` which is not installed yet.'
+            )
         train_transform = amdim_transforms.TransformsSTL10Patches(
             patch_size=patch_size,
             overlap=patch_overlap

pl_bolts/models/self_supervised/amdim/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 AMDIMTrainTransformsCIFAR10:
@@ -78,8 +78,8 @@ class AMDIMEvalTransformsCIFAR10:
     """
 
     def __init__(self):
-        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.'
             )
 
@@ -121,8 +121,8 @@ class AMDIMTrainTransformsSTL10:
     """
 
     def __init__(self, height=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.'
             )
 
@@ -170,8 +170,8 @@ class AMDIMEvalTransformsSTL10(object):
     """
 
     def __init__(self, height=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.'
             )
 
@@ -215,8 +215,8 @@ class AMDIMTrainTransformsImageNet128(object):
     """
 
     def __init__(self, height=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.'
             )
 
@@ -264,8 +264,8 @@ class AMDIMEvalTransformsImageNet128(object):
     """
 
     def __init__(self, height=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.'
             )