upd

Author: AlekseySh

docs/readme/examples_source/extractor/train_val_pl.md

@@ -16,7 +16,7 @@ from oml.miners.inbatch_all_tri import AllTripletsMiner
 from oml.models import ViTExtractor
 from oml.samplers.balance import BalanceSampler
 from oml.utils.download_mock_dataset import download_mock_dataset
-from pytorch_lightning.loggers import NeptuneLogger, TensorBoardLogger, WandbLogger
+from oml.lightning.pipelines.logging import NeptunePipelineLogger, TensorBoardPipelineLogger, WandBPipelineLogger
 
 dataset_root = "mock_dataset/"
 df_train, df_val = download_mock_dataset(dataset_root)
@@ -37,15 +37,15 @@ val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=4)
 metric_callback = MetricValCallback(metric=EmbeddingMetrics(extra_keys=[train_dataset.paths_key,]), log_images=True)
 
 # 1) Logging with Tensorboard
-logger = TensorBoardLogger(".")
+logger = TensorBoardPipelineLogger(".")
 
 # 2) Logging with Neptune
-# logger = NeptuneLogger(api_key="", project="", log_model_checkpoints=False)
+# logger = NeptunePipelineLogger(api_key="", project="", log_model_checkpoints=False)
 
 # 3) Logging with Weights and Biases
 # import os
 # os.environ["WANDB_API_KEY"] = ""
-# logger = WandbLogger(project="test_project", log_model=False)
+# logger = WandBPipelineLogger(project="test_project", log_model=False)
 
 # run
 pl_model = ExtractorModule(extractor, criterion, optimizer)

oml/interfaces/loggers.py

@@ -6,11 +6,13 @@
 from oml.const import TCfg
 
 
-class IPipelineLogger(LightningLogger):
+class IFigureLogger:
     @abstractmethod
-    def log_experiment_info(self, cfg: TCfg) -> None:
+    def log_figure(self, fig: plt.Figure, title: str, idx: int) -> None:
         raise NotImplementedError()
 
+
+class IPipelineLogger(LightningLogger, IFigureLogger):
     @abstractmethod
-    def log_figure(self, fig: plt.Figure, title: str, idx: int) -> None:
+    def log_pipeline_info(self, cfg: TCfg) -> None:
         raise NotImplementedError()

oml/lightning/callbacks/metric.py

@@ -1,3 +1,4 @@
+import warnings
 from math import ceil
 from typing import Any, Optional
 
@@ -9,7 +10,7 @@
 
 from oml.const import LOG_IMAGE_FOLDER
 from oml.ddp.patching import check_loaders_is_patched, patch_dataloader_to_ddp
-from oml.interfaces.loggers import IPipelineLogger
+from oml.interfaces.loggers import IFigureLogger
 from oml.interfaces.metrics import IBasicMetric, IMetricDDP, IMetricVisualisable
 from oml.lightning.modules.ddp import ModuleDDP
 from oml.utils.misc import flatten_dict
@@ -100,12 +101,16 @@ def _log_images(self, pl_module: pl.LightningModule) -> None:
         if not isinstance(self.metric, IMetricVisualisable):
             return
 
+        if not isinstance(pl_module.logger, IFigureLogger):
+            warnings.warn(
+                f"Unexpected logger {pl_module.logger}. Figures have not been saved. "
+                f"Please, use a child of {IFigureLogger}."
+            )
+            return
+
         for fig, metric_log_str in zip(*self.metric.visualize()):
             log_str = f"{LOG_IMAGE_FOLDER}/{metric_log_str}"
-
-            assert isinstance(pl_module.logger, IPipelineLogger)
             pl_module.logger.log_figure(fig=fig, title=log_str, idx=pl_module.current_epoch)
-
             plt.close(fig=fig)
 
     def on_validation_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:

oml/lightning/pipelines/logging.py

@@ -27,7 +27,7 @@ def prepare_tags(cfg: TCfg) -> List[str]:
 
 
 class NeptunePipelineLogger(NeptuneLogger, IPipelineLogger):
-    def log_experiment_info(self, cfg: TCfg) -> None:
+    def log_pipeline_info(self, cfg: TCfg) -> None:
         warnings.warn(
             "Unfortunately, in the case of using Neptune, you may experience that long experiments are "
             "stacked and not responding. It's not an issue on OML's side, so, we cannot fix it."
@@ -57,7 +57,7 @@ def log_figure(self, fig: plt.Figure, title: str, idx: int) -> None:
 
 
 class WandBPipelineLogger(WandbLogger, IPipelineLogger):
-    def log_experiment_info(self, cfg: TCfg) -> None:
+    def log_pipeline_info(self, cfg: TCfg) -> None:
         # this is the optional dependency
         import wandb
 
@@ -90,9 +90,12 @@ def log_figure(self, fig: plt.Figure, title: str, idx: int) -> None:
 
 
 class TensorBoardPipelineLogger(TensorBoardLogger, IPipelineLogger):
-    def log_experiment_info(self, cfg: TCfg) -> None:
+    def log_pipeline_info(self, cfg: TCfg) -> None:
         pass
 
     def log_figure(self, fig: plt.Figure, title: str, idx: int) -> None:
         fig_img = figure_to_nparray(fig)
         self.experiment.add_image(title, np.transpose(fig_img, (2, 0, 1)), idx)
+
+
+__all__ = ["IPipelineLogger", "TensorBoardPipelineLogger", "WandBPipelineLogger", "NeptunePipelineLogger"]

oml/lightning/pipelines/train.py

@@ -72,7 +72,7 @@ def extractor_training_pipeline(cfg: TCfg) -> None:
     pprint(cfg)
 
     logger = parse_logger_from_config(cfg)
-    logger.log_experiment_info(cfg)
+    logger.log_pipeline_info(cfg)
 
     trainer_engine_params = parse_engine_params_from_config(cfg)
     is_ddp = check_is_config_for_ddp(trainer_engine_params)

oml/lightning/pipelines/train_postprocessor.py

@@ -120,7 +120,7 @@ def postprocessor_training_pipeline(cfg: DictConfig) -> None:
     pprint(cfg)
 
     logger = parse_logger_from_config(cfg)
-    logger.log_experiment_info(cfg)
+    logger.log_pipeline_info(cfg)
 
     trainer_engine_params = parse_engine_params_from_config(cfg)
     is_ddp = check_is_config_for_ddp(trainer_engine_params)