diff --git a/README.md b/README.md index ab228321eb..310e946a3a 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ MMEngine is a foundational library for training deep learning models based on Py **Covers mainstream training monitoring platforms** - [TensorBoard](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#tensorboard) | [WandB](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#wandb) | [MLflow](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#mlflow-wip) -- [ClearML](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#clearml) | [Neptune](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#neptune) | [DVCLive](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#dvclive) | [Aim](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#aim) +- [ClearML](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#clearml) | [Neptune](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#neptune) | [DVCLive](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#dvclive) | [Aim](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#aim) | [SwanLab](https://mmengine.readthedocs.io/en/latest/common_usage/visualize_training_log.html#swanlab) ## Installation diff --git a/README_zh-CN.md b/README_zh-CN.md index 7a67cdbf76..d6f743acbc 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -92,7 +92,7 @@ MMEngine 是一个基于 PyTorch 实现的，用于训练深度学习模型的 **覆盖主流的训练监测平台** - [TensorBoard](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#tensorboard) | [WandB](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#wandb) | [MLflow](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#mlflow-wip) -- [ClearML](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#clearml) | [Neptune](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#neptune) | [DVCLive](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#dvclive) | [Aim](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#aim) +- [ClearML](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#clearml) | [Neptune](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#neptune) | [DVCLive](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#dvclive) | [Aim](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#aim) | [SwanLab](https://mmengine.readthedocs.io/zh-cn/latest/common_usage/visualize_training_log.html#swanlab) **兼容主流的训练芯片** diff --git a/docs/en/common_usage/visualize_training_log.md b/docs/en/common_usage/visualize_training_log.md index 239220f3f7..d9fdbedc3c 100644 --- a/docs/en/common_usage/visualize_training_log.md +++ b/docs/en/common_usage/visualize_training_log.md @@ -1,6 +1,6 @@ # Visualize Training Logs -MMEngine integrates experiment management tools such as [TensorBoard](https://www.tensorflow.org/tensorboard), [Weights & Biases (WandB)](https://docs.wandb.ai/), [MLflow](https://mlflow.org/docs/latest/index.html), [ClearML](https://clear.ml/docs/latest/docs), [Neptune](https://docs.neptune.ai/), [DVCLive](https://dvc.org/doc/dvclive) and [Aim](https://aimstack.readthedocs.io/en/latest/overview.html), making it easy to track and visualize metrics like loss and accuracy. +MMEngine integrates experiment management tools such as [TensorBoard](https://www.tensorflow.org/tensorboard), [Weights & Biases (WandB)](https://docs.wandb.ai/), [MLflow](https://mlflow.org/docs/latest/index.html), [ClearML](https://clear.ml/docs/latest/docs), [Neptune](https://docs.neptune.ai/), [DVCLive](https://dvc.org/doc/dvclive), [Aim](https://aimstack.readthedocs.io/en/latest/overview.html) and [SwanLab](https://docs.swanlab.cn/), making it easy to track and visualize metrics like loss and accuracy. Below, we'll show you how to configure an experiment management tool in just one line, based on the example from [15 minutes to get started with MMEngine](../get_started/15_minutes.md). @@ -234,3 +234,48 @@ to launch the Aim UI as shown below. ![image](https://github.com/open-mmlab/mmengine/assets/58739961/2fc6cdd8-1de7-4125-a20a-c95c1a8bdb1b) Initialization configuration parameters are available at [Aim SDK Reference](https://aimstack.readthedocs.io/en/latest/refs/sdk.html#module-aim.sdk.run). + +## SwanLab + +Before using SwanLab, you need to install `swanlab` dependency library. + +```bash +pip install swanlab +swanlab login +``` + +Configure the `visualizer` in the initialization parameters of the Runner, and set `vis_backends` to [SwanLabVisBackend](mmengine.visualization.SwanLabVisBackend). + +```python +runner = Runner( + model=MMResNet50(), + work_dir='./work_dir', + train_dataloader=train_dataloader, + optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)), + train_cfg=dict(by_epoch=True, max_epochs=5, val_interval=1), + val_dataloader=val_dataloader, + val_cfg=dict(), + val_evaluator=dict(type=Accuracy), + visualizer=dict(type='Visualizer', vis_backends=[dict(type='SwanLabVisBackend')]), +) +runner.train() +``` + +You can click on [SwanLabVisBackend API](mmengine.visualization.SwanLabVisBackend) to view the configurable parameters for `SwanLabVisBackend`. For example, the `init_kwargs` parameter will be passed to the [swanlab.init](https://docs.swanlab.cn/en/api/py-init.html) method. + +```python +runner = Runner( + ... + visualizer=dict( + type='Visualizer', + vis_backends=[ + dict( + type='SwanLabVisBackend', + init_kwargs=dict(project='toy-example') + ), + ], + ), + ... +) +runner.train() +``` diff --git a/docs/zh_cn/api/visualization.rst b/docs/zh_cn/api/visualization.rst index 39b2de476f..6dc46a6f2a 100644 --- a/docs/zh_cn/api/visualization.rst +++ b/docs/zh_cn/api/visualization.rst @@ -38,3 +38,4 @@ visualization Backend NeptuneVisBackend DVCLiveVisBackend AimVisBackend + SwanLabVisBackend diff --git a/docs/zh_cn/common_usage/visualize_training_log.md b/docs/zh_cn/common_usage/visualize_training_log.md index ede2152385..2825313603 100644 --- a/docs/zh_cn/common_usage/visualize_training_log.md +++ b/docs/zh_cn/common_usage/visualize_training_log.md @@ -1,6 +1,6 @@ # 可视化训练日志 -MMEngine 集成了 [TensorBoard](https://www.tensorflow.org/tensorboard?hl=zh-cn)、[Weights & Biases (WandB)](https://docs.wandb.ai/)、[MLflow](https://mlflow.org/docs/latest/index.html) 、[ClearML](https://clear.ml/docs/latest/docs)、[Neptune](https://docs.neptune.ai/)、[DVCLive](https://dvc.org/doc/dvclive) 和 [Aim](https://aimstack.readthedocs.io/en/latest/overview.html) 实验管理工具，你可以很方便地跟踪和可视化损失及准确率等指标。 +MMEngine 集成了 [TensorBoard](https://www.tensorflow.org/tensorboard?hl=zh-cn)、[Weights & Biases (WandB)](https://docs.wandb.ai/)、[MLflow](https://mlflow.org/docs/latest/index.html) 、[ClearML](https://clear.ml/docs/latest/docs)、[Neptune](https://docs.neptune.ai/)、[DVCLive](https://dvc.org/doc/dvclive)、[Aim](https://aimstack.readthedocs.io/en/latest/overview.html)和 [SwanLab](https://docs.swanlab.cn/) 实验管理工具，你可以很方便地跟踪和可视化损失及准确率等指标。 下面基于 [15 分钟上手 MMENGINE](../get_started/15_minutes.md) 中的例子介绍如何一行配置实验管理工具。 @@ -234,3 +234,48 @@ aim up ![image](https://github.com/open-mmlab/mmengine/assets/58739961/2fc6cdd8-1de7-4125-a20a-c95c1a8bdb1b) 初始化配置参数可点击 [Aim SDK Reference](https://aimstack.readthedocs.io/en/latest/refs/sdk.html#module-aim.sdk.run) 查询。 + +## SwanLab + +使用 SwanLab 前需先安装 `swanlab` 依赖库。 + +```bash +pip install swanlab +swanlab login +``` + +设置 `Runner` 初始化参数中的 `visualizer`，并将 `vis_backends` 设置为 [SwanLabVisBackend](mmengine.visualization.SwanLabVisBackend)。 + +```python +runner = Runner( + model=MMResNet50(), + work_dir='./work_dir', + train_dataloader=train_dataloader, + optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)), + train_cfg=dict(by_epoch=True, max_epochs=5, val_interval=1), + val_dataloader=val_dataloader, + val_cfg=dict(), + val_evaluator=dict(type=Accuracy), + visualizer=dict(type='Visualizer', vis_backends=[dict(type='SwanLabVisBackend')]), +) +runner.train() +``` + +点击 [SwanLabVisBackend API](mmengine.visualization.SwanLabVisBackend) 查看 `SwanLabVisBackend` 可配置的参数。例如 `init_kwargs` 参数会传给 [swanlab.init](https://docs.swanlab.cn/en/api/py-init.html) 方法。 + +```python +runner = Runner( + ... + visualizer=dict( + type='Visualizer', + vis_backends=[ + dict( + type='SwanLabVisBackend', + init_kwargs=dict(project='toy-example') + ), + ], + ), + ... +) +runner.train() +``` diff --git a/docs/zh_cn/get_started/introduction.md b/docs/zh_cn/get_started/introduction.md index 34fd28630a..4a95f5591c 100644 --- a/docs/zh_cn/get_started/introduction.md +++ b/docs/zh_cn/get_started/introduction.md @@ -22,7 +22,7 @@ MMEngine 是一个基于 PyTorch 实现的，用于训练深度学习模型的 **覆盖主流的训练监测平台** - [TensorBoard](../common_usage/visualize_training_log.md#tensorboard) | [WandB](../common_usage/visualize_training_log.md#wandb) | [MLflow](../common_usage/visualize_training_log.md#mlflow-wip) -- [ClearML](../common_usage/visualize_training_log.md#clearml) | [Neptune](../common_usage/visualize_training_log.md#neptune) | [DVCLive](../common_usage/visualize_training_log.md#dvclive) | [Aim](../common_usage/visualize_training_log.md#aim) +- [ClearML](../common_usage/visualize_training_log.md#clearml) | [Neptune](../common_usage/visualize_training_log.md#neptune) | [DVCLive](../common_usage/visualize_training_log.md#dvclive) | [Aim](../common_usage/visualize_training_log.md#aim) | [SwanLab](../common_usage/visualize_training_log.md#swanlab) **兼容主流的训练芯片** diff --git a/mmengine/visualization/__init__.py b/mmengine/visualization/__init__.py index 8f59452c54..f17925f586 100644 --- a/mmengine/visualization/__init__.py +++ b/mmengine/visualization/__init__.py @@ -1,12 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. from .vis_backend import (AimVisBackend, BaseVisBackend, ClearMLVisBackend, DVCLiveVisBackend, LocalVisBackend, MLflowVisBackend, - NeptuneVisBackend, TensorboardVisBackend, - WandbVisBackend) + NeptuneVisBackend, SwanLabVisBackend, + TensorboardVisBackend, WandbVisBackend) from .visualizer import Visualizer __all__ = [ 'Visualizer', 'BaseVisBackend', 'LocalVisBackend', 'WandbVisBackend', 'TensorboardVisBackend', 'MLflowVisBackend', 'ClearMLVisBackend', - 'NeptuneVisBackend', 'DVCLiveVisBackend', 'AimVisBackend' + 'NeptuneVisBackend', 'DVCLiveVisBackend', 'AimVisBackend', + 'SwanLabVisBackend' ] diff --git a/mmengine/visualization/vis_backend.py b/mmengine/visualization/vis_backend.py index b752ec85a7..e7973ef9b3 100644 --- a/mmengine/visualization/vis_backend.py +++ b/mmengine/visualization/vis_backend.py @@ -1446,3 +1446,157 @@ def close(self) -> None: return self._aim_run.close() + + +@VISBACKENDS.register_module() +class SwanLabVisBackend(BaseVisBackend): + """Swanlab visualization backend class for mmengine. + Examples: + >>> from mmengine.visualization import SwanLabVisBackend + >>> import numpy as np + >>> Swanlab_vis_backend = SwanlabVisBackend() + >>> img=np.random.randint(0, 256, size=(10, 10, 3)) + >>> Swanlab_vis_backend.add_image('img', img) + >>> Swanlab_vis_backend.add_scaler('mAP', 0.6) + >>> Swanlab_vis_backend.add_scalars({'loss': [1, 2, 3],'acc': 0.8}) + >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1]))) + >>> Swanlab_vis_backend.add_config(cfg) + + Args: + save_dir (str, optional): The root directory to save the files + produced by the visualizer. Default used './swanlab' + init_kwargs (dict, optional): Swanlab initialization + input parameters. + See `swanlab.init `_ for + details. Defaults to None. + """ + + def __init__( + self, + save_dir: str = None, + init_kwargs: Optional[dict] = None, + ): + self._save_dir = save_dir + self._env_initialized = False + self._init_kwargs = init_kwargs + + @force_init_env + def experiment(self) -> Any: + """Return the experiment object associated with this visualization + backend. + + The experiment attribute can get the swanlab backend. If you want to + write other data, such as writing a table, you can directly get the + visualization backend through experiment. + """ + return self._swanlab + + def _init_env(self) -> Any: + """Setup env for swanlab.""" + if self._save_dir is not None: + if not os.path.exists(self._save_dir): + os.makedirs(self._save_dir, exist_ok=True) # type: ignore + if self._init_kwargs is None: + self._init_kwargs = {'logdir': self._save_dir} + else: + self._init_kwargs.setdefault('logdir', self._save_dir) + try: + import swanlab + except ImportError: + raise ImportError( + 'Please run "pip install swanlab" to install swanlab') + + swanlab.config['FRAMEWORK'] = 'mmengine' + swanlab.init(**self._init_kwargs) + self._swanlab = swanlab + + @force_init_env + def add_config(self, config: Config, **kwargs) -> None: + """Record the config to swanlab. + + Args: + config (Config): The Config object + """ + + def repack_dict(a, prefix=''): + """Unpack Nested Dictionary func.""" + new_dict = dict() + for key, value in a.items(): + key = str(key) + if isinstance(value, dict): + if prefix != '': + new_dict.update(repack_dict(value, f'{prefix}/{key}')) + else: + new_dict.update(repack_dict(value, key)) + elif isinstance(value, list) or isinstance(value, tuple): + if all(not isinstance(element, dict) for element in value): + new_dict[key] = value + else: + for i, item in enumerate(value): + new_dict.update(repack_dict(item, f'{key}[{i}]')) + elif prefix != '': + new_dict[f'{prefix}/{key}'] = value + else: + new_dict[key] = value + return new_dict + + config_dict = config.to_dict() + self._swanlab.config.update(repack_dict(config_dict)) + + @force_init_env + def add_image(self, + name: str, + image: np.ndarray, + step: int = 0, + **kwargs) -> None: + """Record the image to swanlab. + + Args: + name (str): The image identifier. + image (np.ndarray): The image to be saved. The format + should be RGB. Defaults to None. + step (int): Global step value to record. Defaults to 0. + """ + image = self._swanlab.Image(image) + self._swanlab.log({name: image}, step=step) + + @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 to swanlab. + + Args: + name (str): The scalar identifier. + value (int, float): Value to save. + step (int): Global step value to record. Defaults to 0. + """ + self._swanlab.log({name: value}, step=step) + + @force_init_env + def add_scalars( + self, + scalar_dict: dict, + step: int = 0, + file_path: Optional[str] = None, + **kwargs, + ) -> None: + """Record the scalars' data. + + 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): The scalar's data will be + saved to the `file_path` file at the same time + if the `file_path` parameter is specified. + Defaults to None. + """ + self._swanlab.log(scalar_dict, step=step) + + def close(self) -> None: + """Close an opened swanlab object.""" + if hasattr(self, '_swanlab'): + self._swanlab.finish() diff --git a/tests/test_visualizer/test_vis_backend.py b/tests/test_visualizer/test_vis_backend.py index c991462ef9..1b7572feab 100644 --- a/tests/test_visualizer/test_vis_backend.py +++ b/tests/test_visualizer/test_vis_backend.py @@ -17,7 +17,8 @@ from mmengine.visualization import (AimVisBackend, ClearMLVisBackend, DVCLiveVisBackend, LocalVisBackend, MLflowVisBackend, NeptuneVisBackend, - TensorboardVisBackend, WandbVisBackend) + SwanLabVisBackend, TensorboardVisBackend, + WandbVisBackend) class TestLocalVisBackend: @@ -496,3 +497,50 @@ def test_close(self): aim_vis_backend = AimVisBackend() aim_vis_backend._init_env() aim_vis_backend.close() + + +class TestSwanLabVisBackend: + + def test_init(self): + SwanLabVisBackend('temp_dir') + VISBACKENDS.build(dict(type='SwanLabVisBackend', save_dir='temp_dir')) + + def test_experiment(self): + swanlab_vis_backend = SwanLabVisBackend('temp_dir') + assert swanlab_vis_backend.experiment == swanlab_vis_backend._swanlab + shutil.rmtree('temp_dir') + + def test_add_config(self): + cfg = Config(dict(a=1, b=dict(b1=[0, 1]))) + swanlab_vis_backend = SwanLabVisBackend('temp_dir') + swanlab_vis_backend.add_config(cfg) + shutil.rmtree('temp_dir') + + def test_add_image(self): + image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8) + swanlab_vis_backend = SwanLabVisBackend('temp_dir') + swanlab_vis_backend.add_image('img', image) + swanlab_vis_backend.add_image('img', image) + shutil.rmtree('temp_dir') + + def test_add_scalar(self): + swanlab_vis_backend = SwanLabVisBackend('temp_dir') + swanlab_vis_backend.add_scalar('map', 0.9) + # test append mode + swanlab_vis_backend.add_scalar('map', 0.9) + swanlab_vis_backend.add_scalar('map', 0.95) + shutil.rmtree('temp_dir') + + def test_add_scalars(self): + swanlab_vis_backend = SwanLabVisBackend('temp_dir') + input_dict = {'map': 0.7, 'acc': 0.9} + swanlab_vis_backend.add_scalars(input_dict) + # test append mode + swanlab_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8}) + shutil.rmtree('temp_dir') + + def test_close(self): + swanlab_vis_backend = SwanLabVisBackend('temp_dir') + swanlab_vis_backend._init_env() + swanlab_vis_backend.close() + shutil.rmtree('temp_dir')