Unverified Commit bb06f354 authored by Jerry Jiarui XU's avatar Jerry Jiarui XU Committed by GitHub
Browse files

[Feature] Add hook msg for checkpoint hook (#635)

* [Feature] Add hook msg for ckpt_hook

* [Feature] Add hook msg for ckpt_hook

* add test
parent 9a3f1249
......@@ -60,6 +60,16 @@ class CheckpointHook(Hook):
self.out_dir = runner.work_dir
runner.save_checkpoint(
self.out_dir, save_optimizer=self.save_optimizer, **self.args)
if runner.meta is not None:
if self.by_epoch:
cur_ckpt_filename = self.args.get(
'filename_tmpl', 'epoch_{}.pth').format(runner.epoch + 1)
else:
cur_ckpt_filename = self.args.get(
'filename_tmpl', 'iter_{}.pth').format(runner.iter + 1)
runner.meta.setdefault('hook_msgs', dict())
runner.meta['hook_msgs']['last_ckpt'] = os.path.join(
self.out_dir, cur_ckpt_filename)
# remove other checkpoints
if self.max_keep_ckpts > 0:
if self.by_epoch:
......
......@@ -23,6 +23,32 @@ from mmcv.runner import (CheckpointHook, EMAHook, IterTimerHook,
from mmcv.runner.hooks.lr_updater import CosineRestartLrUpdaterHook
def test_checkpoint_hook():
"""xdoctest -m tests/test_runner/test_hooks.py test_checkpoint_hook."""
# test epoch based runner
loader = DataLoader(torch.ones((5, 2)))
runner = _build_demo_runner('EpochBasedRunner', max_epochs=1)
runner.meta = dict()
checkpointhook = CheckpointHook(interval=1, by_epoch=True)
runner.register_hook(checkpointhook)
runner.run([loader], [('train', 1)])
assert runner.meta['hook_msgs']['last_ckpt'] == osp.join(
runner.work_dir, 'epoch_1.pth')
shutil.rmtree(runner.work_dir)
# test iter based runner
runner = _build_demo_runner(
'IterBasedRunner', max_iters=1, max_epochs=None)
runner.meta = dict()
checkpointhook = CheckpointHook(interval=1, by_epoch=False)
runner.register_hook(checkpointhook)
runner.run([loader], [('train', 1)])
assert runner.meta['hook_msgs']['last_ckpt'] == osp.join(
runner.work_dir, 'iter_1.pth')
shutil.rmtree(runner.work_dir)
def test_ema_hook():
"""xdoctest -m tests/test_hooks.py test_ema_hook."""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment