[CodeCamp2023-Task288] Add NeptuneVisBackend

mmengine/visualization/__init__.py

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

mmengine/visualization/vis_backend.py

@@ -986,3 +986,138 @@ def close(self) -> None:
         for file_path in file_paths:
             self._task.upload_artifact(os.path.basename(file_path), file_path)
         self._task.close()
+
+
+@VISBACKENDS.register_module()
+class NeptuneVisBackend(BaseVisBackend):
+    """Neptune visualization backend class.
+
+    Examples:
+        >>>from mmengine.visualization import NeptuneVisBackend
+        >>> import numpy as np
+        >>> init_kwargs = {'project': 'your_project_name'}
+        >>> neptune_vis_backend = NeptuneVisBackend(init_kwargs=init_kwargs)
+        >>> img = np.random.randint(0, 256, size=(10, 10, 3))
+        >>> neptune_vis_backend.add_image('img', img)
+        >>> neptune_vis_backend.add_scaler('mAP', 0.6)
+        >>> neptune_vis_backend.add_scalars({'loss': 0.1, 'acc': 0.8})
+        >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
+        >>> neptune_vis_backend.add_config(cfg)
+
+    Args:
+        save_dir (str, optional): The root directory to save the files
+            produced by the visualizer. NeptuneVisBackend does
+            not require this argument. Defaults to None.
+        init_kwargs (dict, optional): Neptune initialization parameters.
+            Defaults to None.
+
+            - project (str): Name of a project in a form of
+              `namespace/project_name`. If `project` is not specified,
+              the value of `NEPTUNE_PROJECT` environment variable
+              will be taken.
+            - api_token (str): User’s API token. If api_token is not api_token,
+              the value of `NEPTUNE_API_TOKEN` environment variable will
+              be taken. Note: It is strongly recommended to use
+              `NEPTUNE_API_TOKEN` environment variable rather than
+              placing your API token here.
+
+            If 'project' and 'api_token are not specified in `init_kwargs`
+            The 'mode' will set to 'offline'.
+            See `neptune.init_run
+            <https://docs.neptune.ai/api/neptune/#init_run>`_ for
+            details.
+    """
+
+    def __init__(self, save_dir: Optional[str] = None, init_kwargs: Optional[dict] = None):
+        super().__init__(save_dir)
+        self._init_kwargs = init_kwargs
+
+    def _init_env(self):
+        """Setup env for neptune."""
+        try:
+            import neptune
+        except ImportError:
+            raise ImportError(
+                'Please run "pip install -U neptune" to install neptune')
+        if self._init_kwargs is None:
+            self._init_kwargs = {'mode': 'offline'}
+
+        run = neptune.init_run(**self._init_kwargs)
+        self._neptune = run
+
+    @property  # type: ignore
+    @force_init_env
+    def experiment(self):
+        """Return Neptune object."""
+        return self._neptune
+
+    @force_init_env
+    def add_config(self, config: Config, **kwargs) -> None:
+        """Record the config to neptune.
+
+        Args:
+            config (Config): The Config object
+        """
+        from neptune.types import File
+        self._neptune['config'].upload(File.from_content(config.pretty_text))
+
+    @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 neptune.types import File
+
+        # values in the array need to be in the [0, 1] range
+        img = image.astype(np.float32) / 255.0
+        self._neptune['images'].append(
+            File.as_image(img), name=name, step=step)
+
+    @force_init_env
+    def add_scalar(self,
+                   name: str,
+                   value: Union[int, float],
+                   step: int = 0,
+                   **kwargs) -> None:
+        """Record the scalar.
+
+        Args:
+            name (str): The scalar identifier.
+            value (int, float): Value to save.
+            step (int): Global step value to record. Defaults to 0.
+        """
+        self._neptune[name].append(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.
+        """
+        for k, v in scalar_dict.items():
+            self._neptune[k].append(v, step=step)
+
+    def close(self) -> None:
+        """close an opened object."""
+        if hasattr(self, '_neptune'):
+            self._neptune.stop()

requirements/tests.txt

@@ -4,5 +4,6 @@ dadaptation
 lion-pytorch
 lmdb
 mlflow
+neptune
 parameterized
 pytest

tests/test_visualizer/test_vis_backend.py

@@ -13,8 +13,8 @@
 from mmengine.fileio import load
 from mmengine.registry import VISBACKENDS
 from mmengine.visualization import (ClearMLVisBackend, LocalVisBackend,
-                                    MLflowVisBackend, TensorboardVisBackend,
-                                    WandbVisBackend)
+                                    MLflowVisBackend, NeptuneVisBackend,
+                                    TensorboardVisBackend, WandbVisBackend)
 
 
 class TestLocalVisBackend:
@@ -353,3 +353,51 @@ def test_close(self):
         clearml_vis_backend._init_env()
         clearml_vis_backend.add_config(cfg)
         clearml_vis_backend.close()
+
+
+class TestNeptuneVisBackend:
+
+    def test_init(self):
+        NeptuneVisBackend('temp_dir')
+        VISBACKENDS.build(dict(type='NeptuneVisBackend', save_dir='temp_dir'))
+
+    def test_experiment(self):
+        neptune_vis_backend = NeptuneVisBackend('temp_dir')
+        assert neptune_vis_backend.experiment == neptune_vis_backend._neptune
+        shutil.rmtree('temp_dir')
+
+    def test_add_config(self):
+        cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
+        neptune_vis_backend = NeptuneVisBackend('temp_dir')
+        neptune_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)
+        neptune_vis_backend = NeptuneVisBackend('temp_dir')
+        neptune_vis_backend.add_image('img', image)
+        neptune_vis_backend.add_image('img', image)
+        shutil.rmtree('temp_dir')
+
+    def test_add_scalar(self):
+        neptune_vis_backend = NeptuneVisBackend('temp_dir')
+        neptune_vis_backend.add_scalar('map', 0.9)
+        # test append mode
+        neptune_vis_backend.add_scalar('map', 0.9)
+        neptune_vis_backend.add_scalar('map', 0.95)
+        shutil.rmtree('temp_dir')
+
+    def test_add_scalars(self):
+        neptune_vis_backend = NeptuneVisBackend('temp_dir')
+        input_dict = {'map': 0.7, 'acc': 0.9}
+        neptune_vis_backend.add_scalars(input_dict)
+        # test append mode
+        neptune_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8})
+        neptune_vis_backend.add_scalars({'map': [0.8], 'acc': 0.8})
+        shutil.rmtree('temp_dir')
+
+    def test_close(self):
+        neptune_vis_backend = NeptuneVisBackend('temp_dir')
+        neptune_vis_backend._init_env()
+        neptune_vis_backend.close()
+        shutil.rmtree('temp_dir')