mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] Fix publish multiple checkpoints when using multiple GPUs (#1059) (#1070)

Browse Source
This commit is contained in:
Junwei Zheng 2023-04-12 04:38:48 +02:00 committed by GitHub
parent 9207e84aa0
commit d41906fa15
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mmengine/hooks/checkpoint_hook.py
View File

@ -7,7 +7,7 @@ from math import inf
from pathlib import Path
from typing import Callable, Dict, List, Optional, Sequence, Union
from mmengine.dist import is_main_process
from mmengine.dist import is_main_process, master_only
from mmengine.fileio import FileClient, get_file_backend
from mmengine.logging import print_log
from mmengine.registry import HOOKS
@ -347,6 +347,7 @@ class CheckpointHook(Hook):
for key, best_ckpt in self.best_ckpt_path_dict.items():
self._publish_model(runner, best_ckpt)
@master_only
def _publish_model(self, runner, ckpt_path: str) -> None:
"""Remove unnecessary keys from ckpt_path and save the new checkpoint.