From 68414516aa7008e257e45ff4c7c6ab08958c7a6d Mon Sep 17 00:00:00 2001 From: vugia truong Date: Thu, 1 Jun 2023 18:41:34 +0900 Subject: [PATCH] [Feature] Add vis backend for clearml (#878) (#1091) --- docs/en/api/visualization.rst | 1 + docs/zh_cn/api/visualization.rst | 1 + mmengine/visualization/__init__.py | 7 +- mmengine/visualization/vis_backend.py | 143 +++++++++++++++++++++- requirements/tests.txt | 1 + tests/test_visualizer/test_vis_backend.py | 50 +++++++- 6 files changed, 196 insertions(+), 7 deletions(-) diff --git a/docs/en/api/visualization.rst b/docs/en/api/visualization.rst index 10362a9b..4e1f4aeb 100644 --- a/docs/en/api/visualization.rst +++ b/docs/en/api/visualization.rst @@ -34,3 +34,4 @@ visualization Backend MLflowVisBackend TensorboardVisBackend WandbVisBackend + ClearMLVisBackend diff --git a/docs/zh_cn/api/visualization.rst b/docs/zh_cn/api/visualization.rst index 10362a9b..4e1f4aeb 100644 --- a/docs/zh_cn/api/visualization.rst +++ b/docs/zh_cn/api/visualization.rst @@ -34,3 +34,4 @@ visualization Backend MLflowVisBackend TensorboardVisBackend WandbVisBackend + ClearMLVisBackend diff --git a/mmengine/visualization/__init__.py b/mmengine/visualization/__init__.py index 866fbedd..a0a518e6 100644 --- a/mmengine/visualization/__init__.py +++ b/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' ] diff --git a/mmengine/visualization/vis_backend.py b/mmengine/visualization/vis_backend.py index dbdb28c9..dd9f69ac 100644 --- a/mmengine/visualization/vis_backend.py +++ b/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() diff --git a/requirements/tests.txt b/requirements/tests.txt index 5836a67b..f049bc1e 100644 --- a/requirements/tests.txt +++ b/requirements/tests.txt @@ -1,3 +1,4 @@ +clearml coverage dadaptation lion-pytorch diff --git a/tests/test_visualizer/test_vis_backend.py b/tests/test_visualizer/test_vis_backend.py index 97caf87d..bf8183a2 100644 --- a/tests/test_visualizer/test_vis_backend.py +++ b/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()