Prefer local imports for optional dependencies (#18551)

Author: awaelchli

src/lightning/pytorch/callbacks/rich_model_summary.py

@@ -17,10 +17,6 @@
 from lightning.pytorch.callbacks.progress.rich_progress import _RICH_AVAILABLE
 from lightning.pytorch.utilities.model_summary import get_human_readable_count
 
-if _RICH_AVAILABLE:  # type: ignore[has-type]
-    from rich import get_console
-    from rich.table import Table
-
 
 class RichModelSummary(ModelSummary):
     r"""Generates a summary of all layers in a :class:`~lightning.pytorch.core.module.LightningModule` with `rich text
@@ -74,6 +70,9 @@ def summarize(
         model_size: float,
         **summarize_kwargs: Any,
     ) -> None:
+        from rich import get_console
+        from rich.table import Table
+
         console = get_console()
 
         header_style: str = summarize_kwargs.get("header_style", "bold magenta")

src/lightning/pytorch/core/saving.py

@@ -41,16 +41,10 @@
 from lightning.pytorch.utilities.parsing import AttributeDict, parse_class_init_keys
 from lightning.pytorch.utilities.rank_zero import rank_zero_warn
 
-log = logging.getLogger(__name__)
-
-if _OMEGACONF_AVAILABLE:
-    from omegaconf import OmegaConf
-    from omegaconf.dictconfig import DictConfig
-    from omegaconf.errors import UnsupportedValueType, ValidationError
-
 if TYPE_CHECKING:
     from torch.storage import UntypedStorage
 
+log = logging.getLogger(__name__)
 # the older shall be on the top
 CHECKPOINT_PAST_HPARAMS_KEYS = ("hparams", "module_arguments")  # used in 0.7.6
 
@@ -296,6 +290,9 @@ def load_hparams_from_yaml(config_yaml: _PATH, use_omegaconf: bool = True) -> Di
         hparams = yaml.full_load(fp)
 
     if _OMEGACONF_AVAILABLE and use_omegaconf:
+        from omegaconf import OmegaConf
+        from omegaconf.errors import UnsupportedValueType, ValidationError
+
         with contextlib.suppress(UnsupportedValueType, ValidationError):
             return OmegaConf.create(hparams)
     return hparams
@@ -322,6 +319,10 @@ def save_hparams_to_yaml(config_yaml: _PATH, hparams: Union[dict, Namespace], us
 
     # saving with OmegaConf objects
     if _OMEGACONF_AVAILABLE and use_omegaconf:
+        from omegaconf import OmegaConf
+        from omegaconf.dictconfig import DictConfig
+        from omegaconf.errors import UnsupportedValueType, ValidationError
+
         # deepcopy: hparams from user shouldn't be resolved
         hparams = deepcopy(hparams)
         hparams = apply_to_collection(hparams, DictConfig, OmegaConf.to_container, resolve=True)

src/lightning/pytorch/demos/mnist_datamodule.py

@@ -28,9 +28,6 @@
 from lightning.pytorch import LightningDataModule
 from lightning.pytorch.utilities.imports import _TORCHVISION_AVAILABLE
 
-if _TORCHVISION_AVAILABLE:
-    from torchvision import transforms as transform_lib
-
 _DATASETS_PATH = "./data"
 
 
@@ -244,11 +241,14 @@ def test_dataloader(self) -> DataLoader:
     def default_transforms(self) -> Optional[Callable]:
         if not _TORCHVISION_AVAILABLE:
             return None
+
+        from torchvision import transforms
+
         if self.normalize:
-            mnist_transforms = transform_lib.Compose(
-                [transform_lib.ToTensor(), transform_lib.Normalize(mean=(0.5,), std=(0.5,))]
+            mnist_transforms = transforms.Compose(
+                [transforms.ToTensor(), transforms.Normalize(mean=(0.5,), std=(0.5,))]
             )
         else:
-            mnist_transforms = transform_lib.ToTensor()
+            mnist_transforms = transforms.ToTensor()
 
         return mnist_transforms

src/lightning/pytorch/loggers/tensorboard.py

@@ -37,9 +37,6 @@
 
 log = logging.getLogger(__name__)
 
-if _OMEGACONF_AVAILABLE:
-    from omegaconf import Container, OmegaConf
-
 # Skip doctests if requirements aren't available
 if not (_TENSORBOARD_AVAILABLE or _TENSORBOARDX_AVAILABLE):
     __doctest_skip__ = ["TensorBoardLogger", "TensorBoardLogger.*"]
@@ -178,6 +175,9 @@ def log_hyperparams(  # type: ignore[override]
             metrics: Dictionary with metric names as keys and measured quantities as values
 
         """
+        if _OMEGACONF_AVAILABLE:
+            from omegaconf import Container, OmegaConf
+
         params = _convert_params(params)
 
         # store params to output

src/lightning/pytorch/loops/evaluation_loop.py

@@ -44,10 +44,6 @@
 from lightning.pytorch.utilities.model_helpers import is_overridden
 from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
 
-if _RICH_AVAILABLE:
-    from rich import get_console
-    from rich.table import Column, Table
-
 
 class _EvaluationLoop(_Loop):
     """Top-level loop where validation/testing starts."""
@@ -531,6 +527,9 @@ def _print_results(results: List[_OUT_DICT], stage: str) -> None:
             table_headers.insert(0, f"{stage} Metric".capitalize())
 
             if _RICH_AVAILABLE:
+                from rich import get_console
+                from rich.table import Column, Table
+
                 columns = [Column(h, justify="center", style="magenta", width=max_length) for h in table_headers]
                 columns[0].style = "cyan"
 

src/lightning/pytorch/strategies/ddp.py

@@ -14,7 +14,7 @@
 import logging
 from contextlib import nullcontext
 from datetime import timedelta
-from typing import Any, Callable, Dict, List, Literal, Optional, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, TYPE_CHECKING, Union
 
 import torch
 import torch.distributed
@@ -47,7 +47,7 @@
 from lightning.pytorch.utilities.exceptions import _augment_message
 from lightning.pytorch.utilities.rank_zero import rank_zero_deprecation, rank_zero_info, rank_zero_only
 
-if torch.distributed.is_available():
+if TYPE_CHECKING:
     from torch.distributed.algorithms.model_averaging.averagers import ModelAverager
 
 log = logging.getLogger(__name__)

src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

@@ -35,11 +35,7 @@
 from lightning.pytorch.utilities.migration.utils import _pl_migrate_checkpoint
 from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_warn
 
-if _OMEGACONF_AVAILABLE:
-    from omegaconf import Container
-
-
-log: logging.Logger = logging.getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 class _CheckpointConnector:
@@ -464,6 +460,9 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
                 checkpoint[prec_plugin.__class__.__qualname__] = prec_plugin_state_dict
             prec_plugin.on_save_checkpoint(checkpoint)
 
+        if _OMEGACONF_AVAILABLE:
+            from omegaconf import Container
+
         # dump hyper-parameters
         for obj in (model, datamodule):
             if obj and obj.hparams:

src/lightning/pytorch/tuner/lr_finder.py

@@ -36,10 +36,11 @@
 else:
     from tqdm import tqdm
 
-_MATPLOTLIB_AVAILABLE = RequirementCache("matplotlib")
-if TYPE_CHECKING and _MATPLOTLIB_AVAILABLE:
+if TYPE_CHECKING:
     import matplotlib.pyplot as plt
     from matplotlib.axes import Axes
+
+_MATPLOTLIB_AVAILABLE = RequirementCache("matplotlib")
 log = logging.getLogger(__name__)
 
 

src/lightning/pytorch/utilities/deepspeed.py

@@ -23,13 +23,6 @@
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.strategies.deepspeed import _DEEPSPEED_AVAILABLE
 
-if _DEEPSPEED_AVAILABLE:
-    from deepspeed.utils.zero_to_fp32 import (
-        get_fp32_state_dict_from_zero_checkpoint,
-        get_model_state_file,
-        get_optim_files,
-    )
-
 CPU_DEVICE = torch.device("cpu")
 
 
@@ -75,6 +68,14 @@ def convert_zero_checkpoint_to_fp32_state_dict(
             "lightning_model.pt"
         )
     """
+    if not _DEEPSPEED_AVAILABLE:
+        raise ModuleNotFoundError(str(_DEEPSPEED_AVAILABLE))
+
+    from deepspeed.utils.zero_to_fp32 import (
+        get_fp32_state_dict_from_zero_checkpoint,
+        get_model_state_file,
+        get_optim_files,
+    )
 
     state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)