From a8005d14f9505e2adbcd44170d6cf974ed602a47 Mon Sep 17 00:00:00 2001 From: Marchematics Date: Sun, 30 Aug 2026 22:39:58 +0800 Subject: [PATCH 1/3] fix: save images on all processes --- mmengine/visualization/visualizer.py | 1 - tests/test_visualizer/test_visualizer.py | 12 ++++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) 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_visualizer/test_visualizer.py b/tests/test_visualizer/test_visualizer.py index e4ababc637..a772ceb766 100644 --- a/tests/test_visualizer/test_visualizer.py +++ b/tests/test_visualizer/test_visualizer.py @@ -554,6 +554,18 @@ 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') + + 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), From 85b078d12a3ce92a58777ff2edf27eb0915bb3ed Mon Sep 17 00:00:00 2001 From: Marchematics Date: Tue, 1 Sep 2026 03:59:26 +0800 Subject: [PATCH 2/3] Avoid rank collisions in visualizer image files Signed-off-by: Marchematics --- mmengine/visualization/vis_backend.py | 7 ++++++- tests/test_visualizer/test_vis_backend.py | 12 ++++++++++++ tests/test_visualizer/test_visualizer.py | 1 + 3 files changed, 19 insertions(+), 1 deletion(-) 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/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 a772ceb766..1569e72f37 100644 --- a/tests/test_visualizer/test_visualizer.py +++ b/tests/test_visualizer/test_visualizer.py @@ -560,6 +560,7 @@ def test_add_image_on_non_master_process(self, is_main_process): 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']: From d76767e53de6835acca2ea7e6ecfe724a8847205 Mon Sep 17 00:00:00 2001 From: Marchematics Date: Tue, 1 Sep 2026 08:31:10 +0800 Subject: [PATCH 3/3] Fix duplicate mock in checkpoint test Signed-off-by: Marchematics --- tests/test_hooks/test_checkpoint_hook.py | 1 - 1 file changed, 1 deletion(-) 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, \