[Feature] Add vis backend for clearml (#878) (#1091)

docs/en/api/visualization.rst

@@ -34,3 +34,4 @@ visualization Backend
    MLflowVisBackend
    TensorboardVisBackend
    WandbVisBackend
+   ClearMLVisBackend

docs/zh_cn/api/visualization.rst

@@ -34,3 +34,4 @@ visualization Backend
    MLflowVisBackend
    TensorboardVisBackend
    WandbVisBackend
+   ClearMLVisBackend

mmengine/visualization/__init__.py

@@ -1,9 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from .vis_backend import (BaseVisBackend, LocalVisBackend, MLflowVisBackend,
-                          TensorboardVisBackend, WandbVisBackend)
+from .vis_backend import (BaseVisBackend, ClearMLVisBackend, LocalVisBackend,
+                          MLflowVisBackend, TensorboardVisBackend,
+                          WandbVisBackend)
 from .visualizer import Visualizer
 
 __all__ = [
     'Visualizer', 'BaseVisBackend', 'LocalVisBackend', 'WandbVisBackend',
-    'TensorboardVisBackend', 'MLflowVisBackend'
+    'TensorboardVisBackend', 'MLflowVisBackend', 'ClearMLVisBackend'
 ]

mmengine/visualization/vis_backend.py

@@ -7,7 +7,7 @@ import os.path as osp
 import warnings
 from abc import ABCMeta, abstractmethod
 from collections.abc import MutableMapping
-from typing import Any, Callable, Optional, Sequence, Union
+from typing import Any, Callable, List, Optional, Sequence, Union
 
 import cv2
 import numpy as np
@@ -830,3 +830,144 @@ class MLflowVisBackend(BaseVisBackend):
             else:
                 items[new_key] = v
         return items
+
+
+@VISBACKENDS.register_module()
+class ClearMLVisBackend(BaseVisBackend):
+    """Clearml visualization backend class. It requires `clearml`_ to be
+    installed.
+
+    Examples:
+        >>> from mmengine.visualization import ClearMLVisBackend
+        >>> from mmengine import Config
+        >>> import numpy as np
+        >>> vis_backend = ClearMLVisBackend(save_dir='temp_dir')
+        >>> img = np.random.randint(0, 256, size=(10, 10, 3))
+        >>> vis_backend.add_image('img.png', img)
+        >>> vis_backend.add_scalar('mAP', 0.6)
+        >>> vis_backend.add_scalars({'loss': 0.1,'acc':0.8})
+        >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
+        >>> vis_backend.add_config(cfg)
+
+    Args:
+        save_dir (str, optional): Useless parameter. Just for
+            interface unification. Defaults to None.
+        init_kwargs (dict, optional): A dict contains the arguments of
+            ``clearml.Task.init`` . See `taskinit`_  for more details.
+            Defaults to None
+        artifact_suffix (Tuple[str] or str): The artifact suffix.
+            Defaults to ('.py', 'pth').
+
+    .. _clearml:
+        https://clear.ml/docs/latest/docs/
+
+    .. _taskinit:
+        https://clear.ml/docs/latest/docs/references/sdk/task/#taskinit
+    """
+
+    def __init__(self,
+                 save_dir: Optional[str] = None,
+                 init_kwargs: Optional[dict] = None,
+                 artifact_suffix: SUFFIX_TYPE = ('.py', '.pth')):
+        super().__init__(save_dir)  # type: ignore
+        self._init_kwargs = init_kwargs
+        self._artifact_suffix = artifact_suffix
+
+    def _init_env(self) -> None:
+        try:
+            import clearml
+        except ImportError:
+            raise ImportError(
+                'Please run "pip install clearml" to install clearml')
+
+        task_kwargs = self._init_kwargs or {}
+        self._clearml = clearml
+        self._task = self._clearml.Task.init(**task_kwargs)
+        self._logger = self._task.get_logger()
+
+    @property  # type: ignore
+    @force_init_env
+    def experiment(self):
+        """Return clearml object."""
+        return self._clearml
+
+    @force_init_env
+    def add_config(self, config: Config, **kwargs) -> None:
+        """Record the config to clearml.
+
+        Args:
+            config (Config): The Config object
+        """
+        self.cfg = config
+        self._task.connect_configuration(vars(config))
+
+    @force_init_env
+    def add_image(self,
+                  name: str,
+                  image: np.ndarray,
+                  step: int = 0,
+                  **kwargs) -> None:
+        """Record the image to clearml.
+
+        Args:
+            name (str): The image identifier.
+            image (np.ndarray): The image to be saved. The format
+                should be RGB.
+            step (int): Global step value to record. Defaults to 0.
+        """
+        self._logger.report_image(
+            title=name, series=name, iteration=step, image=image)
+
+    @force_init_env
+    def add_scalar(self,
+                   name: str,
+                   value: Union[int, float, torch.Tensor, np.ndarray],
+                   step: int = 0,
+                   **kwargs) -> None:
+        """Record the scalar data to clearml.
+
+        Args:
+            name (str): The scalar identifier.
+            value (int, float, torch.Tensor, np.ndarray): Value to save.
+            step (int): Global step value to record. Defaults to 0.
+        """
+        self._logger.report_scalar(
+            title=name, series=name, value=value, iteration=step)
+
+    @force_init_env
+    def add_scalars(self,
+                    scalar_dict: dict,
+                    step: int = 0,
+                    file_path: Optional[str] = None,
+                    **kwargs) -> None:
+        """Record the scalar's data to clearml.
+
+        Args:
+            scalar_dict (dict): Key-value pair storing the tag and
+                corresponding values.
+            step (int): Global step value to record. Defaults to 0.
+            file_path (str, optional): Useless parameter. Just for
+                interface unification. Defaults to None.
+        """
+        assert 'step' not in scalar_dict, 'Please set it directly ' \
+                                          'through the step parameter'
+        for key, value in scalar_dict.items():
+            self._logger.report_scalar(
+                title=key, series=key, value=value, iteration=step)
+
+    def close(self) -> None:
+        """Close the clearml."""
+        if not hasattr(self, '_clearml'):
+            return
+
+        file_paths: List[str] = list()
+        if (hasattr(self, 'cfg')
+                and osp.isdir(getattr(self.cfg, 'work_dir', ''))):
+            for filename in scandir(self.cfg.work_dir, self._artifact_suffix,
+                                    False):
+                file_path = osp.join(self.cfg.work_dir, filename)
+                file_paths.append(file_path)
+
+        for file_path in file_paths:
+            self._task.upload_artifact(os.path.basename(file_path), file_path)
+        self._task.close()

requirements/tests.txt

@@ -1,3 +1,4 @@
+clearml
 coverage
 dadaptation
 lion-pytorch

tests/test_visualizer/test_vis_backend.py

@@ -3,7 +3,7 @@ import os
 import shutil
 import sys
 import warnings
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
@@ -12,8 +12,9 @@ import torch
 from mmengine import Config
 from mmengine.fileio import load
 from mmengine.registry import VISBACKENDS
-from mmengine.visualization import (LocalVisBackend, MLflowVisBackend,
-                                    TensorboardVisBackend, WandbVisBackend)
+from mmengine.visualization import (ClearMLVisBackend, LocalVisBackend,
+                                    MLflowVisBackend, TensorboardVisBackend,
+                                    WandbVisBackend)
 
 
 class TestLocalVisBackend:
@@ -285,3 +286,46 @@ class TestMLflowVisBackend:
         mlflow_vis_backend.add_config(cfg)
         mlflow_vis_backend.close()
         shutil.rmtree('temp_dir')
+
+
+@patch.dict(sys.modules, {'clearml': MagicMock()})
+class TestClearMLVisBackend:
+
+    def test_init(self):
+        ClearMLVisBackend('temp_dir')
+        VISBACKENDS.build(dict(type='ClearMLVisBackend', save_dir='temp_dir'))
+
+    def test_experiment(self):
+        clearml_vis_backend = ClearMLVisBackend('temp_dir')
+        assert clearml_vis_backend.experiment == clearml_vis_backend._clearml
+
+    def test_add_config(self):
+        cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
+        clearml_vis_backend = ClearMLVisBackend('temp_dir')
+        clearml_vis_backend.add_config(cfg)
+
+    def test_add_image(self):
+        image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
+        clearml_vis_backend = ClearMLVisBackend('temp_dir')
+        clearml_vis_backend.add_image('img.png', image)
+
+    def test_add_scalar(self):
+        clearml_vis_backend = ClearMLVisBackend('temp_dir')
+        clearml_vis_backend.add_scalar('map', 0.9)
+        # test append mode
+        clearml_vis_backend.add_scalar('map', 0.9)
+        clearml_vis_backend.add_scalar('map', 0.95)
+
+    def test_add_scalars(self):
+        clearml_vis_backend = ClearMLVisBackend('temp_dir')
+        input_dict = {'map': 0.7, 'acc': 0.9}
+        clearml_vis_backend.add_scalars(input_dict)
+        # test append mode
+        clearml_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8})
+
+    def test_close(self):
+        cfg = Config(dict(work_dir='temp_dir'))
+        clearml_vis_backend = ClearMLVisBackend('temp_dir')
+        clearml_vis_backend._init_env()
+        clearml_vis_backend.add_config(cfg)
+        clearml_vis_backend.close()