Add clearml logger

Add clearml logger

Author: Fissium

README.md

@@ -392,7 +392,13 @@ 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 oml.lightning.pipelines.logging import NeptunePipelineLogger, TensorBoardPipelineLogger, WandBPipelineLogger, MLFlowPipelineLogger
+from oml.lightning.pipelines.logging import (
+    ClearMLPipelineLogger,
+    MLFlowPipelineLogger,
+    NeptunePipelineLogger,
+    TensorBoardPipelineLogger,
+    WandBPipelineLogger,
+)
 
 dataset_root = "mock_dataset/"
 df_train, df_val = download_mock_dataset(dataset_root)
@@ -426,6 +432,9 @@ logger = TensorBoardPipelineLogger(".")
 # 4) Logging with MLFlow locally
 # logger = MLFlowPipelineLogger(experiment_name="exp", tracking_uri="file:./ml-runs")
 
+# 5) Logging with ClearML
+# logger = ClearMLPipelineLogger(project_name="exp", task_name="test")
+
 # run
 pl_model = ExtractorModule(extractor, criterion, optimizer)
 trainer = pl.Trainer(max_epochs=3, callbacks=[metric_callback], num_sanity_val_steps=0, logger=logger)

ci/requirements_optional.txt

@@ -3,4 +3,5 @@ jupyter>=1.0.0
 neptune>=1.0.0, <1.10.1
 wandb>=0.15.4
 mlflow>=2.0.0
+clearml>=1.5.0
 

docs/readme/examples_source/extractor/train_val_pl.md

@@ -16,7 +16,13 @@ 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 oml.lightning.pipelines.logging import NeptunePipelineLogger, TensorBoardPipelineLogger, WandBPipelineLogger, MLFlowPipelineLogger
+from oml.lightning.pipelines.logging import (
+    ClearMLPipelineLogger,
+    MLFlowPipelineLogger,
+    NeptunePipelineLogger,
+    TensorBoardPipelineLogger,
+    WandBPipelineLogger,
+)
 
 dataset_root = "mock_dataset/"
 df_train, df_val = download_mock_dataset(dataset_root)
@@ -50,6 +56,9 @@ logger = TensorBoardPipelineLogger(".")
 # 4) Logging with MLFlow locally
 # logger = MLFlowPipelineLogger(experiment_name="exp", tracking_uri="file:./ml-runs")
 
+# 5) Logging with ClearML
+# logger = ClearMLPipelineLogger(project_name="exp", task_name="test")
+
 # run
 pl_model = ExtractorModule(extractor, criterion, optimizer)
 trainer = pl.Trainer(max_epochs=3, callbacks=[metric_callback], num_sanity_val_steps=0, logger=logger)

docs/source/oml/logging.rst

@@ -60,6 +60,19 @@ There are several loggers integrated with Pipelines. You can also `use your cust
               tracking_uri: "file:./ml-runs"  # another way: export MLFLOW_TRACKING_URI=file:./ml-runs
       ...
 
+* ClearML
+
+  .. code-block:: yaml
+
+      ...
+      logger:
+          name: clearml
+          args:
+              project_name: "test_project"
+              task_name: "test"
+              offline_mode: False # if True logging is directed to a local dir
+      ...
+
 
 An example of logging via Neptune in the
 `feature extractor <https://github.com/OML-Team/open-metric-learning/tree/main/pipelines/features_extraction>`_
@@ -121,6 +134,7 @@ Take a look at the following example:
 `Training + Validation [Lightning and logging] <https://open-metric-learning.readthedocs.io/en/latest/feature_extraction/python_examples.html>`_.
 It shows how to use each of: `Tensorboard <https://pytorch.org/docs/stable/tensorboard.html>`_,
 `MLFlow <mlflow.org>`_,
+`ClearML <https://clear.ml/>`_,
 `Neptune <https://neptune.ai/>`_ or
 `WandB <https://wandb.ai/site>`_.
 

oml/configs/logger/clearml.yaml

@@ -0,0 +1,4 @@
+name: clearml
+args:
+  project_name: "test_project"
+  task_name: "test"

oml/lightning/pipelines/logging.py

@@ -1,10 +1,14 @@
 import warnings
+from argparse import Namespace
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Mapping, Optional, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
+from lightning_fabric.utilities.logger import _flatten_dict
+from lightning_fabric.utilities.rank_zero import rank_zero_only
 from pytorch_lightning.loggers import (
+    Logger,
     MLFlowLogger,
     NeptuneLogger,
     TensorBoardLogger,
@@ -31,6 +35,66 @@ def prepare_tags(cfg: TCfg) -> List[str]:
     return tags
 
 
+class ClearMLLogger(Logger):
+    def __init__(self, **kwargs: Any):
+        try:
+            from clearml import Task
+        except ImportError as e:
+            raise ModuleNotFoundError(
+                "This contrib module requires clearml to be installed. "
+                "You may install clearml using: \n pip install clearml \n"
+            ) from e
+
+        experiment_kwargs = {
+            k: v for k, v in kwargs.items() if k not in ("project_name", "task_name", "task_type", "offline_mode")
+        }
+
+        if kwargs.get("offline_mode", False):
+            Task.set_offline(offline_mode=True)
+            warnings.warn("ClearMLSaver: running in offline mode")
+
+        # Try to retrieve current the ClearML Task before trying to create a new one
+        self.task = Task.current_task()
+        if self.task is None:
+            self.task = Task.init(
+                project_name=kwargs.get("project_name"),
+                task_name=kwargs.get("task_name"),
+                task_type=kwargs.get("task_type", Task.TaskTypes.training),
+                **experiment_kwargs,
+            )
+
+        self.logger = self.task.get_logger()
+
+    @property
+    def name(self) -> str:
+        return "ClearMLLogger"
+
+    @property
+    def version(self) -> Union[int, str]:
+        return self.task.id
+
+    @rank_zero_only
+    def finalize(self, status: str) -> None:
+        self.logger.flush()
+
+    @rank_zero_only
+    def log_hyperparams(self, params: Optional[Union[Dict[str, Any], Namespace]]) -> None:
+        if isinstance(params, Namespace):
+            params = vars(params)
+
+        if params is None:
+            params = {}
+        params = _flatten_dict(params)
+
+        self.task.connect(params)
+
+    @rank_zero_only
+    def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None) -> None:
+        assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"  # type: ignore
+        for k, v in metrics.items():
+            self.logger.report_scalar(title=k, series=k, iteration=step, value=v)
+
+
 class NeptunePipelineLogger(NeptuneLogger, IPipelineLogger):
     def log_pipeline_info(self, cfg: TCfg) -> None:
         warnings.warn(
@@ -132,10 +196,44 @@ def log_figure(self, fig: plt.Figure, title: str, idx: int) -> None:
         self.experiment.log_figure(figure=fig, artifact_file=f"{title}.png", run_id=self.run_id)
 
 
+class ClearMLPipelineLogger(ClearMLLogger, IPipelineLogger):
+    def log_pipeline_info(self, cfg: TCfg) -> None:
+        # log config
+        self.log_hyperparams(prepare_config_to_logging(cfg))
+
+        # log tags
+        self.task.add_tags(prepare_tags(cfg))
+
+        # log transforms as files
+        names_files = save_transforms_as_files(cfg)
+        if names_files:
+            for name, transforms_file in names_files:
+                self.task.upload_artifact(name=name, artifact_object=transforms_file)
+
+        # log code
+        self.task.upload_artifact(name="code", artifact_object=OML_PATH)
+
+        # log dataframe
+        self.task.upload_artifact(
+            name="dataset",
+            artifact_object=str(Path(cfg["dataset_root"]) / cfg["dataframe_name"]),
+        )
+
+    def log_figure(self, fig: plt.Figure, title: str, idx: int) -> None:
+        self.logger.report_matplotlib_figure(
+            title=title,
+            series="",
+            figure=fig,
+            iteration=idx,
+            report_image=True,
+        )
+
+
 __all__ = [
     "IPipelineLogger",
     "TensorBoardPipelineLogger",
     "WandBPipelineLogger",
     "NeptunePipelineLogger",
     "MLFlowPipelineLogger",
+    "ClearMLPipelineLogger",
 ]

oml/registry/loggers.py

@@ -4,6 +4,7 @@
 from oml.const import TCfg
 from oml.interfaces.loggers import IPipelineLogger
 from oml.lightning.pipelines.logging import (
+    ClearMLPipelineLogger,
     MLFlowPipelineLogger,
     NeptunePipelineLogger,
     TensorBoardPipelineLogger,
@@ -16,6 +17,7 @@
     "neptune": NeptunePipelineLogger,
     "tensorboard": TensorBoardPipelineLogger,
     "mlflow": MLFlowPipelineLogger,
+    "clearml": ClearMLPipelineLogger,
 }
 
 CLOUD_TOKEN_NAMES = {"wandb": "WANDB_API_KEY", "neptune": "NEPTUNE_API_TOKEN"}

tests/test_imports.py

@@ -9,7 +9,7 @@
 
 from oml.const import PROJECT_ROOT
 
-LIBS_TO_IGNORE = ["torch_xla", "pytorch_grad_cam", "wandb", "neptune", "IPython"]
+LIBS_TO_IGNORE = ["torch_xla", "pytorch_grad_cam", "wandb", "neptune", "clearml", "IPython"]
 
 NEED_TO_TEST_NOTEBOOKS = True
 

tests/test_runs/test_pipelines/configs/train_arcface_with_categories.yaml

@@ -51,17 +51,19 @@ metric_args:
   return_only_overall_category: True
   visualize_only_overall_category: True
 
-log_images: False
+log_images: True
 
 metric_for_checkpointing: OVERALL/cmc/1
 
 max_epochs: 2
 valid_period: 1
 
 logger:
-  name: tensorboard
+  name: clearml
   args:
-    save_dir: "."
+    project_name: "test_project"
+    task_name: "test"
+    offline_mode: True
 
 tags:
   - mock

tests/test_runs/test_pipelines/test_pipelines.py

@@ -80,6 +80,7 @@ def test_train_with_categories(accelerator: str, devices: int) -> None:
 
 
 @pytest.mark.long
+@pytest.mark.needs_optional_dependency
 @pytest.mark.parametrize("accelerator, devices", accelerator_devices_pairs())
 def test_train_arcface_with_categories(accelerator: str, devices: int) -> None:
     run("train_arcface_with_categories.py", accelerator, devices)