Make the lightning dependency optional

mart/__init__.py

@@ -1,11 +1,14 @@
-import importlib
+import importlib.metadata
 
 from mart import attack as attack
-from mart import datamodules as datamodules
-from mart import models as models
 from mart import nn as nn
 from mart import optim as optim
 from mart import transforms as transforms
 from mart import utils as utils
+from mart.utils.imports import _HAS_LIGHTNING
+
+if _HAS_LIGHTNING:
+    from mart import datamodules as datamodules
+    from mart import models as models
 
 __version__ = importlib.metadata.version(__package__ or __name__)

mart/attack/__init__.py

@@ -1,4 +1,4 @@
-from .adversary import *
+from ..utils.imports import _HAS_LIGHTNING
 from .adversary_wrapper import *
 from .composer import *
 from .enforcer import *
@@ -8,3 +8,6 @@
 from .objective import *
 from .perturber import *
 from .projector import *
+
+if _HAS_LIGHTNING:
+    from .adversary import *

mart/attack/initializer/base.py

@@ -10,10 +10,6 @@
 
 import torch
 
-from mart.utils import pylogger
-
-logger = pylogger.get_pylogger(__name__)
-
 
 class Initializer:
     """Initializer base class."""

mart/attack/initializer/vision.py

@@ -4,14 +4,15 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
+import logging
+
 import torch
 import torchvision
 import torchvision.transforms.functional as F
 
-from ...utils import pylogger
 from .base import Initializer
 
-logger = pylogger.get_pylogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 class Image(Initializer):

mart/attack/perturber.py

@@ -9,7 +9,6 @@
 from typing import TYPE_CHECKING, Callable, Iterable, Sequence
 
 import torch
-from lightning.pytorch.utilities.exceptions import MisconfigurationException
 
 from .projector import Projector
 
@@ -85,21 +84,19 @@ def create_from_tensor(tensor):
 
     def named_parameters(self, *args, **kwargs):
         if self.perturbation is None:
-            raise MisconfigurationException("You need to call configure_perturbation before fit.")
+            raise RuntimeError("You need to call configure_perturbation before fit.")
 
         return super().named_parameters(*args, **kwargs)
 
     def parameters(self, *args, **kwargs):
         if self.perturbation is None:
-            raise MisconfigurationException("You need to call configure_perturbation before fit.")
+            raise RuntimeError("You need to call configure_perturbation before fit.")
 
         return super().parameters(*args, **kwargs)
 
     def forward(self, **batch):
         if self.perturbation is None:
-            raise MisconfigurationException(
-                "You need to call the configure_perturbation before forward."
-            )
+            raise RuntimeError("You need to call the configure_perturbation before forward.")
 
         self.projector_(self.perturbation, **batch)
         # We need to register the hook at every forward pass.

mart/callbacks/__init__.py

@@ -1,3 +1,4 @@
+# All Lightning callbacks dependent on lightning, so we don't import mart.callbacks by default.
 from ..utils.imports import _HAS_TORCHVISION
 from .adversary_connector import *
 from .eval_mode import *

mart/utils/__init__.py

@@ -1,11 +1,12 @@
+# Only import components without external dependency.
 from .adapters import *
-from .config import *
-from .imports import _HAS_TORCHVISION
+from .imports import _HAS_LIGHTNING
 from .monkey_patch import *
-from .pylogger import *
-from .rich_utils import *
+from .optimization import *
 from .silent import *
 from .utils import *
 
-if _HAS_TORCHVISION:
-    from .export import *
+if _HAS_LIGHTNING:
+    from .lightning import *
+    from .pylogger import *
+    from .rich_utils import *

mart/utils/imports.py

@@ -4,17 +4,19 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
+import logging
 from importlib.util import find_spec
 
-from .pylogger import get_pylogger
-
-logger = get_pylogger(__name__)
+# Avoid importing .pylogger when checking imports before running other code.
+logger = logging.getLogger(__name__)
 
 
 def has(module_name):
     module = find_spec(module_name)
     if module is None:
-        logger.warn(f"{module_name} is not installed, so some features in MART are unavailable.")
+        logger.warning(
+            f"{module_name} is not installed, so some features in MART are unavailable."
+        )
         return False
     else:
         return True
@@ -25,3 +27,4 @@ def has(module_name):
 _HAS_TORCHVISION = has("torchvision")
 _HAS_TIMM = has("timm")
 _HAS_PYCOCOTOOLS = has("pycocotools")
+_HAS_LIGHTNING = has("lightning")