Refactor pl_bolts.callbacks

pl_bolts/callbacks/__init__.py

@@ -1,7 +1,21 @@
 """
 Collection of PyTorchLightning callbacks
 """
+from pl_bolts.callbacks.byol_updates import BYOLMAWeightUpdate  # noqa: F401
 from pl_bolts.callbacks.data_monitor import ModuleDataMonitor, TrainingDataMonitor  # noqa: F401
 from pl_bolts.callbacks.printing import PrintTableMetricsCallback  # noqa: F401
+from pl_bolts.callbacks.ssl_online import SSLOnlineEvaluator  # noqa: F401
 from pl_bolts.callbacks.variational import LatentDimInterpolator  # noqa: F401
-from pl_bolts.callbacks.vision import TensorboardGenerativeModelImageSampler  # noqa: F401
+from pl_bolts.callbacks.vision.confused_logit import ConfusedLogitCallback  # noqa: F401
+from pl_bolts.callbacks.vision.image_generation import TensorboardGenerativeModelImageSampler  # noqa: F401
+
+__all__ = [
+    "BYOLMAWeightUpdate",
+    "ModuleDataMonitor",
+    "TrainingDataMonitor",
+    "PrintTableMetricsCallback",
+    "SSLOnlineEvaluator",
+    "LatentDimInterpolator",
+    "ConfusedLogitCallback",
+    "TensorboardGenerativeModelImageSampler",
+]

pl_bolts/callbacks/data_monitor.py

@@ -10,9 +10,13 @@
 from torch import Tensor
 from torch.utils.hooks import RemovableHandle
 
-try:
+from pl_bolts.utils import _WANDB_AVAILABLE
+from pl_bolts.utils.warnings import warn_missing_pkg
+
+if _WANDB_AVAILABLE:
     import wandb
-except ModuleNotFoundError:
+else:  # pragma: no cover
+    warn_missing_pkg("wandb")
     wandb = None
 
 
@@ -87,6 +91,11 @@ def log_histogram(self, tensor: Tensor, name: str) -> None:
             )
 
         if isinstance(logger, WandbLogger):
+            if not _WANDB_AVAILABLE:  # pragma: no cover
+                raise ModuleNotFoundError(
+                    "You want to use `wandb` which is not installed yet."
+                )
+
             logger.experiment.log(
                 data={name: wandb.Histogram(tensor)}, commit=False,
             )

pl_bolts/callbacks/ssl_online.py

@@ -6,7 +6,7 @@
 from torch.nn import functional as F
 
 
-class SSLOnlineEvaluator(Callback):  # pragma: no-cover
+class SSLOnlineEvaluator(Callback):  # pragma: no cover
     """
     Attaches a MLP for fine-tuning using the standard self-supervised protocol.
 

pl_bolts/callbacks/variational.py

@@ -4,12 +4,13 @@
 import torch
 from pytorch_lightning.callbacks import Callback
 
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-try:
+if _TORCHVISION_AVAILABLE:
     import torchvision
-except ModuleNotFoundError:
-    warn_missing_pkg('torchvision')  # pragma: no-cover
+else:  # pragma: no cover
+    warn_missing_pkg("torchvision")
 
 
 class LatentDimInterpolator(Callback):
@@ -44,6 +45,11 @@ def __init__(
             num_samples: default 2
             normalize: default True (change image to (0, 1) range)
         """
+        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
+            raise ModuleNotFoundError(
+                "You want to use `torchvision` which is not installed yet."
+            )
+
         super().__init__()
         self.interpolate_epoch_interval = interpolate_epoch_interval
         self.range_start = range_start

pl_bolts/callbacks/vision/__init__.py

@@ -1,2 +0,0 @@
-from pl_bolts.callbacks.vision.confused_logit import ConfusedLogitCallback  # noqa: F401
-from pl_bolts.callbacks.vision.image_generation import TensorboardGenerativeModelImageSampler  # noqa: F401

pl_bolts/utils/__init__.py

@@ -8,3 +8,4 @@
 _SKLEARN_AVAILABLE: bool = _module_available("sklearn")
 _PIL_AVAILABLE: bool = _module_available("PIL")
 _OPENCV_AVAILABLE: bool = _module_available("cv2")
+_WANDB_AVAILABLE: bool = _module_available("wandb")