[CodeCamp2023-668] Add AimVisBackend

docs/en/api/visualization.rst

@@ -37,3 +37,4 @@ visualization Backend
    ClearMLVisBackend
    NeptuneVisBackend
    DVCLiveVisBackend
+   AimVisBackend

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/) and [DVCLive](https://dvc.org/doc/dvclive), 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) and [Aim](https://aimstack.readthedocs.io/en/latest/overview.html), 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).
 
@@ -190,3 +190,47 @@ Open the `report.html` file under `work_dir_dvc`, and you will see the visualiza
 You can also configure a VSCode extension of [DVC](https://marketplace.visualstudio.com/items?itemName=Iterative.dvc) to visualize the training process.
 
 More initialization configuration parameters are available at [DVCLive API Reference](https://dvc.org/doc/dvclive/live).
+
+## Aim
+
+Before using Aim, you need to install `aim` dependency library.
+
+```bash
+pip install aim
+```
+
+Configure the `Runner` in the initialization parameters of the Runner, and set `vis_backends` to [AimVisBackend](mmengine.visualization.AimVisBackend).
+
+```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='AimVisBackend')]),
+)
+runner.train()
+```
+
+In the terminal, use the following command,
+
+```bash
+aim up
+```
+
+or in the Jupyter Notebook, use the following command,
+
+```bash
+%load_ext aim
+%aim up
+```
+
+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).

docs/zh_cn/api/visualization.rst

@@ -37,3 +37,4 @@ visualization Backend
    ClearMLVisBackend
    NeptuneVisBackend
    DVCLiveVisBackend
+   AimVisBackend

docs/zh_cn/common_usage/visualize_training_log.md

@@ -1,8 +1,8 @@
 # 可视化训练日志
 
-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) 实验管理工具，你可以很方便地跟踪和可视化损失及准确率等指标。
+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) 实验管理工具，你可以很方便地跟踪和可视化损失及准确率等指标。
 
-下面基于[15 分钟上手 MMENGINE](../get_started/15_minutes.md)中的例子介绍如何一行配置实验管理工具。
+下面基于 [15 分钟上手 MMENGINE](../get_started/15_minutes.md) 中的例子介绍如何一行配置实验管理工具。
 
 ## TensorBoard
 
@@ -190,3 +190,47 @@ runner.train()
 你还可以安装 VSCode 扩展 [DVC](https://marketplace.visualstudio.com/items?itemName=Iterative.dvc) 进行可视化。
 
 更多初始化配置参数可点击 [DVCLive API Reference](https://dvc.org/doc/dvclive/live) 查询。
+
+## Aim
+
+使用 Aim 前需先安装依赖库 `aim`。
+
+```bash
+pip install aim
+```
+
+设置 `Runner` 初始化参数中的 `visualizer`，并将 `vis_backends` 设置为 [AimVisBackend](mmengine.visualization.AimVisBackend)。
+
+```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='AimVisBackend')]),
+)
+runner.train()
+```
+
+在终端中输入
+
+```bash
+aim up
+```
+
+或者在 Jupyter Notebook 中输入
+
+```bash
+%load_ext aim
+%aim up
+```
+
+即可启动 Aim UI，界面如下图所示。
+
+![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) 查询。

mmengine/visualization/__init__.py

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

mmengine/visualization/vis_backend.py

@@ -1309,3 +1309,132 @@ def _to_dvc_paramlike(self,
             return value.item()
         else:
             return value
+
+
+@VISBACKENDS.register_module()
+class AimVisBackend(BaseVisBackend):
+    """Aim visualization backend class.
+
+    Examples:
+        >>> from mmengine.visualization import AimVisBackend
+        >>> import numpy as np
+        >>> aim_vis_backend = AimVisBackend(save_dir='temp_dir')
+        >>> img=np.random.randint(0, 256, size=(10, 10, 3))
+        >>> aim_vis_backend.add_image('img', img)
+        >>> aim_vis_backend.add_scalar('mAP', 0.6)
+        >>> aim_vis_backend.add_scalars({'loss': 0.1, 'acc': 0.8})
+        >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
+        >>> aim_vis_backend.add_config(cfg)
+
+    Note:
+        1. `New in version 0.8.5.`
+        2. Refer to
+           `Github issue <https://github.com/aimhubio/aim/issues/2064>`_ ,
+           Aim is not unable to be install on Windows for now.
+
+    Args:
+        save_dir (str, optional): The root directory to save the files
+            produced by the visualizer.
+        init_kwargs (dict, optional): Aim initialization parameters. See
+            `Aim <https://aimstack.readthedocs.io/en/latest/refs/sdk.html>`_
+            for details. Defaults to None.
+    """
+
+    def __init__(self,
+                 save_dir: Optional[str] = None,
+                 init_kwargs: Optional[dict] = None):
+        super().__init__(save_dir)  # type:ignore
+        self._init_kwargs = init_kwargs
+
+    def _init_env(self):
+        """Setup env for Aim."""
+        try:
+            from aim import Run
+        except ImportError:
+            raise ImportError('Please run "pip install aim" to install aim')
+
+        from datetime import datetime
+
+        if self._save_dir is not None:
+            path_list = os.path.normpath(self._save_dir).split(os.sep)
+            exp_name = f'{path_list[-2]}_{path_list[-1]}'
+        else:
+            exp_name = datetime.now().strftime('%Y%m%d_%H%M%S')
+
+        if self._init_kwargs is None:
+            self._init_kwargs = {}
+        self._init_kwargs.setdefault('experiment', exp_name)
+        self._aim_run = Run(**self._init_kwargs)
+
+    @property  # type: ignore
+    @force_init_env
+    def experiment(self):
+        """Return Aim object."""
+        return self._aim_run
+
+    @force_init_env
+    def add_config(self, config, **kwargs) -> None:
+        """Record the config to Aim.
+
+        Args:
+            config (Config): The Config object
+        """
+        self._aim_run['hparams'] = dict(config)
+
+    @force_init_env
+    def add_image(self,
+                  name: str,
+                  image: np.ndarray,
+                  step: int = 0,
+                  **kwargs) -> None:
+        """Record the image.
+
+        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.
+        """
+        from aim import Image
+        self._aim_run.track(name=name, value=Image(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 data to Aim.
+
+        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._aim_run.track(name=name, value=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 scalar's data to wandb.
+
+        Args:
+            scalar_dict (dict): Key-value pair storing the tag and
+                corresponding values.
+            step (int): Useless parameter. Wandb does not
+                need this parameter. Defaults to 0.
+            file_path (str, optional): Useless parameter. Just for
+                interface unification. Defaults to None.
+        """
+        for key, value in scalar_dict.items():
+            self._aim_run.track(name=key, value=value, step=step)
+
+    def close(self) -> None:
+        """Close the Aim."""
+        if not hasattr(self, '_aim_run'):
+            return
+
+        self._aim_run.close()

requirements/tests.txt

@@ -1,3 +1,4 @@
+aim
 clearml
 coverage
 dadaptation

tests/test_visualizer/test_vis_backend.py

@@ -14,10 +14,10 @@
 from mmengine.fileio import load
 from mmengine.registry import VISBACKENDS
 from mmengine.utils import digit_version
-from mmengine.visualization import (ClearMLVisBackend, DVCLiveVisBackend,
-                                    LocalVisBackend, MLflowVisBackend,
-                                    NeptuneVisBackend, TensorboardVisBackend,
-                                    WandbVisBackend)
+from mmengine.visualization import (AimVisBackend, ClearMLVisBackend,
+                                    DVCLiveVisBackend, LocalVisBackend,
+                                    MLflowVisBackend, NeptuneVisBackend,
+                                    TensorboardVisBackend, WandbVisBackend)
 
 
 class TestLocalVisBackend:
@@ -445,3 +445,44 @@ def test_close(self):
         dvclive_vis_backend.add_config(cfg)
         dvclive_vis_backend.close()
         shutil.rmtree('temp_dir')
+
+
+@pytest.mark.skipif(
+    platform.system() == 'Windows',
+    reason='Aim does not support Windows for now.')
+class TestAimVisBackend:
+
+    def test_init(self):
+        AimVisBackend()
+        VISBACKENDS.build(dict(type='AimVisBackend'))
+
+    def test_experiment(self):
+        aim_vis_backend = AimVisBackend()
+        assert aim_vis_backend.experiment == aim_vis_backend._aim_run
+
+    def test_add_config(self):
+        cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
+        aim_vis_backend = AimVisBackend()
+        aim_vis_backend.add_config(cfg)
+
+    def test_add_image(self):
+        image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
+        aim_vis_backend = AimVisBackend()
+        aim_vis_backend.add_image('img', image)
+        aim_vis_backend.add_image('img', image, step=1)
+
+    def test_add_scalar(self):
+        aim_vis_backend = AimVisBackend()
+        aim_vis_backend.add_scalar('map', 0.9)
+        aim_vis_backend.add_scalar('map', 0.9, step=1)
+        aim_vis_backend.add_scalar('map', 0.95, step=2)
+
+    def test_add_scalars(self):
+        aim_vis_backend = AimVisBackend()
+        input_dict = {'map': 0.7, 'acc': 0.9}
+        aim_vis_backend.add_scalars(input_dict)
+
+    def test_close(self):
+        aim_vis_backend = AimVisBackend()
+        aim_vis_backend._init_env()
+        aim_vis_backend.close()