[Feat] add swanlab logger

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
 

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)
 
 **兼容主流的训练芯片**
 

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()
+```

docs/zh_cn/api/visualization.rst

@@ -38,3 +38,4 @@ visualization Backend
    NeptuneVisBackend
    DVCLiveVisBackend
    AimVisBackend
+   SwanLabVisBackend

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()
+```

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)
 
 **兼容主流的训练芯片**
 

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'
 ]

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 <NEED UPDATE>`_ 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()

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')