diff --git a/docs/en/api/visualization.rst b/docs/en/api/visualization.rst index 5265dc6edb..10362a9beb 100644 --- a/docs/en/api/visualization.rst +++ b/docs/en/api/visualization.rst @@ -31,5 +31,6 @@ visualization Backend BaseVisBackend LocalVisBackend + MLflowVisBackend TensorboardVisBackend WandbVisBackend diff --git a/docs/zh_cn/api/visualization.rst b/docs/zh_cn/api/visualization.rst index 5265dc6edb..10362a9beb 100644 --- a/docs/zh_cn/api/visualization.rst +++ b/docs/zh_cn/api/visualization.rst @@ -31,5 +31,6 @@ visualization Backend BaseVisBackend LocalVisBackend + MLflowVisBackend TensorboardVisBackend WandbVisBackend diff --git a/mmengine/hooks/logger_hook.py b/mmengine/hooks/logger_hook.py index 04e5f07294..ddf0227ea9 100644 --- a/mmengine/hooks/logger_hook.py +++ b/mmengine/hooks/logger_hook.py @@ -288,6 +288,9 @@ def after_run(self, runner) -> None: runner (Runner): The runner of the training/testing/validation process. """ + # close the visualizer + runner.visualizer.close() + # copy or upload logs to self.out_dir if self.out_dir is None: return diff --git a/mmengine/visualization/__init__.py b/mmengine/visualization/__init__.py index 6c8b0bb546..866fbedda9 100644 --- a/mmengine/visualization/__init__.py +++ b/mmengine/visualization/__init__.py @@ -1,9 +1,9 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .vis_backend import (BaseVisBackend, LocalVisBackend, +from .vis_backend import (BaseVisBackend, LocalVisBackend, MLflowVisBackend, TensorboardVisBackend, WandbVisBackend) from .visualizer import Visualizer __all__ = [ 'Visualizer', 'BaseVisBackend', 'LocalVisBackend', 'WandbVisBackend', - 'TensorboardVisBackend' + 'TensorboardVisBackend', 'MLflowVisBackend' ] diff --git a/mmengine/visualization/vis_backend.py b/mmengine/visualization/vis_backend.py index 5af5ec7056..e8daeeab95 100644 --- a/mmengine/visualization/vis_backend.py +++ b/mmengine/visualization/vis_backend.py @@ -6,6 +6,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 import cv2 @@ -14,8 +15,10 @@ from mmengine.config import Config from mmengine.fileio import dump -from mmengine.logging import print_log +from mmengine.hooks.logger_hook import SUFFIX_TYPE +from mmengine.logging import MMLogger, print_log from mmengine.registry import VISBACKENDS +from mmengine.utils import scandir from mmengine.utils.dl_utils import TORCH_VERSION @@ -613,3 +616,203 @@ def close(self): """close an opened tensorboard object.""" if hasattr(self, '_tensorboard'): self._tensorboard.close() + + +@VISBACKENDS.register_module() +class MLflowVisBackend(BaseVisBackend): + """MLflow visualization backend class. + + It can write images, config, scalars, etc. to a + mlflow file. + + Examples: + >>> from mmengine.visualization import MLflowVisBackend + >>> from mmengine import Config + >>> import numpy as np + >>> vis_backend = MLflowVisBackend(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): The root directory to save the files + produced by the backend. + exp_name (str, optional): The experiment name. Default to None. + run_name (str, optional): The run name. Default to None. + tags (dict, optional): The tags to be added to the experiment. + Default to None. + params (dict, optional): The params to be added to the experiment. + Default to None. + tracking_uri (str, optional): The tracking uri. Default to None. + artifact_suffix (Tuple[str] or str, optional): The artifact suffix. + Default to ('.json', '.log', '.py', 'yaml'). + """ + + def __init__(self, + save_dir: str, + exp_name: Optional[str] = None, + run_name: Optional[str] = None, + tags: Optional[dict] = None, + params: Optional[dict] = None, + tracking_uri: Optional[str] = None, + artifact_suffix: SUFFIX_TYPE = ('.json', '.log', '.py', + 'yaml')): + super().__init__(save_dir) + self._exp_name = exp_name + self._run_name = run_name + self._tags = tags + self._params = params + self._tracking_uri = tracking_uri + self._artifact_suffix = artifact_suffix + + def _init_env(self): + """Setup env for MLflow.""" + if not os.path.exists(self._save_dir): + os.makedirs(self._save_dir, exist_ok=True) # type: ignore + + try: + import mlflow + except ImportError: + raise ImportError( + 'Please run "pip install mlflow" to install mlflow' + ) # type: ignore + self._mlflow = mlflow + + # when mlflow is imported, a default logger is created. + # at this time, the default logger's stream is None + # so the stream is reopened only when the stream is None + # or the stream is closed + logger = MMLogger.get_current_instance() + for handler in logger.handlers: + if handler.stream is None or handler.stream.closed: + handler.stream = open(handler.baseFilename, 'a') + + if self._tracking_uri is not None: + logger.warning( + 'Please make sure that the mlflow server is running.') + self._mlflow.set_tracking_uri(self._tracking_uri) + else: + if os.name == 'nt': + file_url = f'file:\\{os.path.abspath(self._save_dir)}' + else: + file_url = f'file://{os.path.abspath(self._save_dir)}' + self._mlflow.set_tracking_uri(file_url) + + self._exp_name = self._exp_name or 'Default' + + if self._mlflow.get_experiment_by_name(self._exp_name) is None: + self._mlflow.create_experiment(self._exp_name) + + self._mlflow.set_experiment(self._exp_name) + + if self._run_name is not None: + self._mlflow.set_tag('mlflow.runName', self._run_name) + if self._tags is not None: + self._mlflow.set_tags(self._tags) + if self._params is not None: + self._mlflow.log_params(self._params) + + @property # type: ignore + @force_init_env + def experiment(self): + """Return MLflow object.""" + return self._mlflow + + @force_init_env + def add_config(self, config: Config, **kwargs) -> None: + """Record the config to mlflow. + + Args: + config (Config): The Config object + """ + self.cfg = config + self._mlflow.log_params(self._flatten(self.cfg)) + self._mlflow.log_text(self.cfg.pretty_text, 'config.py') + + @force_init_env + def add_image(self, + name: str, + image: np.ndarray, + step: int = 0, + **kwargs) -> None: + """Record the image to mlflow. + + 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. Default to 0. + """ + self._mlflow.log_image(image, name) + + @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 mlflow. + + Args: + name (str): The scalar identifier. + value (int, float, torch.Tensor, np.ndarray): Value to save. + step (int): Global step value to record. Default to 0. + """ + self._mlflow.log_metric(name, value, 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 mlflow. + + Args: + scalar_dict (dict): Key-value pair storing the tag and + corresponding values. + step (int): Global step value to record. Default to 0. + file_path (str, optional): Useless parameter. Just for + interface unification. Default to None. + """ + assert isinstance(scalar_dict, dict) + assert 'step' not in scalar_dict, 'Please set it directly ' \ + 'through the step parameter' + self._mlflow.log_metrics(scalar_dict, step) + + def close(self) -> None: + """Close the mlflow.""" + file_paths = dict() + for filename in scandir(self.cfg.work_dir, self._artifact_suffix, + True): + file_path = osp.join(self.cfg.work_dir, filename) + relative_path = os.path.relpath(file_path, self.cfg.work_dir) + dir_path = os.path.dirname(relative_path) + file_paths[file_path] = dir_path + + for file_path, dir_path in file_paths.items(): + self._mlflow.log_artifact(file_path, dir_path) + + if hasattr(self, '_mlflow'): + self._mlflow.end_run() + + def _flatten(self, d, parent_key='', sep='.') -> dict: + """Flatten the dict.""" + items = dict() + for k, v in d.items(): + new_key = parent_key + sep + k if parent_key else k + if isinstance(v, MutableMapping): + items.update(self._flatten(v, new_key, sep=sep)) + elif isinstance(v, list): + if any(isinstance(x, dict) for x in v): + for i, x in enumerate(v): + items.update( + self._flatten(x, new_key + sep + str(i), sep=sep)) + else: + items[new_key] = v + else: + items[new_key] = v + return items diff --git a/requirements/tests.txt b/requirements/tests.txt index 5228ea3976..5836a67b8e 100644 --- a/requirements/tests.txt +++ b/requirements/tests.txt @@ -2,5 +2,6 @@ coverage dadaptation lion-pytorch lmdb +mlflow parameterized pytest diff --git a/tests/test_visualizer/test_vis_backend.py b/tests/test_visualizer/test_vis_backend.py index caf4cb7424..97caf87d7e 100644 --- a/tests/test_visualizer/test_vis_backend.py +++ b/tests/test_visualizer/test_vis_backend.py @@ -12,8 +12,8 @@ from mmengine import Config from mmengine.fileio import load from mmengine.registry import VISBACKENDS -from mmengine.visualization import (LocalVisBackend, TensorboardVisBackend, - WandbVisBackend) +from mmengine.visualization import (LocalVisBackend, MLflowVisBackend, + TensorboardVisBackend, WandbVisBackend) class TestLocalVisBackend: @@ -242,3 +242,46 @@ def test_close(self): wandb_vis_backend._init_env() wandb_vis_backend.close() shutil.rmtree('temp_dir') + + +class TestMLflowVisBackend: + + def test_init(self): + MLflowVisBackend('temp_dir') + VISBACKENDS.build(dict(type='MLflowVisBackend', save_dir='temp_dir')) + + def test_experiment(self): + mlflow_vis_backend = MLflowVisBackend('temp_dir') + assert mlflow_vis_backend.experiment == mlflow_vis_backend._mlflow + + def test_add_config(self): + cfg = Config(dict(a=1, b=dict(b1=[0, 1]))) + mlflow_vis_backend = MLflowVisBackend('temp_dir') + mlflow_vis_backend.add_config(cfg) + + def test_add_image(self): + image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8) + mlflow_vis_backend = MLflowVisBackend('temp_dir') + mlflow_vis_backend.add_image('img.png', image) + + def test_add_scalar(self): + mlflow_vis_backend = MLflowVisBackend('temp_dir') + mlflow_vis_backend.add_scalar('map', 0.9) + # test append mode + mlflow_vis_backend.add_scalar('map', 0.9) + mlflow_vis_backend.add_scalar('map', 0.95) + + def test_add_scalars(self): + mlflow_vis_backend = MLflowVisBackend('temp_dir') + input_dict = {'map': 0.7, 'acc': 0.9} + mlflow_vis_backend.add_scalars(input_dict) + # test append mode + mlflow_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8}) + + def test_close(self): + cfg = Config(dict(work_dir='temp_dir')) + mlflow_vis_backend = MLflowVisBackend('temp_dir') + mlflow_vis_backend._init_env() + mlflow_vis_backend.add_config(cfg) + mlflow_vis_backend.close() + shutil.rmtree('temp_dir')