Import lightning.pytorch instead of pytorch_lightning

Author: lucmos

src/nn_core/callbacks.py

@@ -2,9 +2,9 @@
 from pathlib import Path
 from typing import Any, Dict, Optional
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
+from lightning.pytorch import Callback, Trainer
 from omegaconf import DictConfig
-from pytorch_lightning import Callback, Trainer
 
 from nn_core.model_logging import NNLogger
 from nn_core.resume import parse_restore
@@ -57,4 +57,3 @@ def on_save_checkpoint(
         metadata = getattr(trainer.datamodule, "metadata", None)
         if metadata is not None:
             checkpoint[METADATA_KEY] = metadata
-

src/nn_core/common/utils.py

@@ -6,8 +6,8 @@
 import dotenv
 import numpy as np
 from hydra.core.hydra_config import HydraConfig
+from lightning.pytorch import seed_everything
 from omegaconf import DictConfig
-from pytorch_lightning import seed_everything
 from rich.prompt import Prompt
 
 pylogger = logging.getLogger(__name__)

src/nn_core/model_logging.py

@@ -6,11 +6,10 @@
 from typing import Any, Dict, Optional, Union
 
 import hydra
-import pytorch_lightning
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.loggers.logger import Logger
 from omegaconf import DictConfig, OmegaConf
-from pytorch_lightning import LightningModule, Trainer
-from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.loggers.logger import Logger
 
 from nn_core.common import PROJECT_ROOT
 
@@ -115,7 +114,7 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
 
         This method logs metrics as as soon as it received them. If you want to aggregate
         metrics for one specific `step`, use the
-        :meth:`~pytorch_lightning.loggers.base.Logger.agg_and_log_metrics` method.
+        :meth:`~lightning.pytorch.loggers.base.Logger.agg_and_log_metrics` method.
 
         Args:
             metrics: Dictionary with metric names as keys and measured quantities as values
@@ -167,7 +166,7 @@ def run_dir(self) -> str:
 
     def log_configuration(
         self,
-        model: pytorch_lightning.LightningModule,
+        model: LightningModule,
         cfg: Union[Dict[str, Any], argparse.Namespace, DictConfig] = None,
         *args,
         **kwargs,

src/nn_core/serialization.py

@@ -10,10 +10,10 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import torch
-from pytorch_lightning.core.saving import _load_state
-from pytorch_lightning.plugins import TorchCheckpointIO
+from lightning.pytorch.core.saving import _load_state
+from lightning.pytorch.plugins import TorchCheckpointIO
 
 METADATA_KEY: str = "metadata"