Save visualizer images on all processes

mmengine/visualization/vis_backend.py

@@ -15,6 +15,7 @@
 import torch
 
 from mmengine.config import Config, ConfigDict
+from mmengine.dist import get_rank
 from mmengine.fileio import dump
 from mmengine.hooks.logger_hook import SUFFIX_TYPE
 from mmengine.logging import MMLogger, print_log
@@ -256,7 +257,11 @@ def add_image(self,
         assert image.dtype == np.uint8
         drawn_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
         os.makedirs(self._img_save_dir, exist_ok=True)
-        save_file_name = f'{name}_{step}.png'
+        save_file_name = f'{name}_{step}'
+        rank = get_rank()
+        if rank > 0:
+            save_file_name += f'_rank{rank}'
+        save_file_name += '.png'
         cv2.imwrite(osp.join(self._img_save_dir, save_file_name), drawn_image)
 
     @force_init_env

mmengine/visualization/visualizer.py

@@ -1082,7 +1082,6 @@ def add_graph(self, model: torch.nn.Module, data_batch: Sequence[dict],
         for vis_backend in self._vis_backends.values():
             vis_backend.add_graph(model, data_batch, **kwargs)
 
-    @master_only
     def add_image(self, name: str, image: np.ndarray, step: int = 0) -> None:
         """Record the image.
 

tests/test_hooks/test_checkpoint_hook.py

@@ -336,7 +336,6 @@ def test_after_val_epoch(self):
         petrel_client = MagicMock()
         for by_epoch, cfg in [(True, self.epoch_based_cfg),
                               (False, self.iter_based_cfg)]:
-            isfile = MagicMock(return_value=True)
             self.clear_work_dir()
             with patch.dict(sys.modules, {'petrel_client': petrel_client}), \
                  patch('mmengine.fileio.backends.PetrelBackend.put') as put_mock, \

tests/test_visualizer/test_vis_backend.py

@@ -62,6 +62,18 @@ def test_add_image(self):
             os.path.join(local_vis_backend._img_save_dir, 'img_2.png'))
         shutil.rmtree('temp_dir')
 
+    @patch('mmengine.visualization.vis_backend.get_rank', return_value=1)
+    def test_add_image_on_non_master_process(self, get_rank):
+        image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
+        local_vis_backend = LocalVisBackend('temp_dir')
+
+        local_vis_backend.add_image('img', image)
+
+        assert os.path.exists(
+            os.path.join(local_vis_backend._img_save_dir, 'img_0_rank1.png'))
+        get_rank.assert_called()
+        shutil.rmtree('temp_dir')
+
     def test_add_scalar(self):
         local_vis_backend = LocalVisBackend('temp_dir')
         local_vis_backend.add_scalar('map', 0.9)

tests/test_visualizer/test_visualizer.py

@@ -554,6 +554,19 @@ def test_add_image(self):
         for name in ['mock1', 'mock2']:
             assert visualizer.get_backend(name)._add_image is True
 
+    @patch('mmengine.dist.utils.is_main_process', return_value=False)
+    def test_add_image_on_non_master_process(self, is_main_process):
+        visualizer = Visualizer(
+            vis_backends=copy.deepcopy(self.vis_backend_cfg),
+            save_dir='temp_dir')
+
+        is_main_process.reset_mock()
+        visualizer.add_image('img', self.image)
+
+        for name in ['mock1', 'mock2']:
+            assert visualizer._vis_backends[name]._add_image is True
+        is_main_process.assert_not_called()
+
     def test_add_scalar(self):
         visualizer = Visualizer(
             vis_backends=copy.deepcopy(self.vis_backend_cfg),