diff --git a/mmengine/visualization/vis_backend.py b/mmengine/visualization/vis_backend.py index b752ec85a7..e776f1c8fa 100644 --- a/mmengine/visualization/vis_backend.py +++ b/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 diff --git a/mmengine/visualization/visualizer.py b/mmengine/visualization/visualizer.py index 6979395aca..5725645f89 100644 --- a/mmengine/visualization/visualizer.py +++ b/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. diff --git a/tests/test_hooks/test_checkpoint_hook.py b/tests/test_hooks/test_checkpoint_hook.py index 13914341f7..b078a25fcb 100644 --- a/tests/test_hooks/test_checkpoint_hook.py +++ b/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, \ diff --git a/tests/test_visualizer/test_vis_backend.py b/tests/test_visualizer/test_vis_backend.py index c991462ef9..82d50e434b 100644 --- a/tests/test_visualizer/test_vis_backend.py +++ b/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) diff --git a/tests/test_visualizer/test_visualizer.py b/tests/test_visualizer/test_visualizer.py index e4ababc637..1569e72f37 100644 --- a/tests/test_visualizer/test_visualizer.py +++ b/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),