| | |
| | import os.path as osp |
| | import warnings |
| |
|
| | from annotator.uniformer.mmcv.fileio import FileClient |
| | from ..dist_utils import allreduce_params, master_only |
| | from .hook import HOOKS, Hook |
| |
|
| |
|
| | @HOOKS.register_module() |
| | class CheckpointHook(Hook): |
| | """Save checkpoints periodically. |
| | |
| | Args: |
| | interval (int): The saving period. If ``by_epoch=True``, interval |
| | indicates epochs, otherwise it indicates iterations. |
| | Default: -1, which means "never". |
| | by_epoch (bool): Saving checkpoints by epoch or by iteration. |
| | Default: True. |
| | save_optimizer (bool): Whether to save optimizer state_dict in the |
| | checkpoint. It is usually used for resuming experiments. |
| | Default: True. |
| | out_dir (str, optional): The root directory to save checkpoints. If not |
| | specified, ``runner.work_dir`` will be used by default. If |
| | specified, the ``out_dir`` will be the concatenation of ``out_dir`` |
| | and the last level directory of ``runner.work_dir``. |
| | `Changed in version 1.3.16.` |
| | max_keep_ckpts (int, optional): The maximum checkpoints to keep. |
| | In some cases we want only the latest few checkpoints and would |
| | like to delete old ones to save the disk space. |
| | Default: -1, which means unlimited. |
| | save_last (bool, optional): Whether to force the last checkpoint to be |
| | saved regardless of interval. Default: True. |
| | sync_buffer (bool, optional): Whether to synchronize buffers in |
| | different gpus. Default: False. |
| | file_client_args (dict, optional): Arguments to instantiate a |
| | FileClient. See :class:`mmcv.fileio.FileClient` for details. |
| | Default: None. |
| | `New in version 1.3.16.` |
| | |
| | .. warning:: |
| | Before v1.3.16, the ``out_dir`` argument indicates the path where the |
| | checkpoint is stored. However, since v1.3.16, ``out_dir`` indicates the |
| | root directory and the final path to save checkpoint is the |
| | concatenation of ``out_dir`` and the last level directory of |
| | ``runner.work_dir``. Suppose the value of ``out_dir`` is "/path/of/A" |
| | and the value of ``runner.work_dir`` is "/path/of/B", then the final |
| | path will be "/path/of/A/B". |
| | """ |
| |
|
| | def __init__(self, |
| | interval=-1, |
| | by_epoch=True, |
| | save_optimizer=True, |
| | out_dir=None, |
| | max_keep_ckpts=-1, |
| | save_last=True, |
| | sync_buffer=False, |
| | file_client_args=None, |
| | **kwargs): |
| | self.interval = interval |
| | self.by_epoch = by_epoch |
| | self.save_optimizer = save_optimizer |
| | self.out_dir = out_dir |
| | self.max_keep_ckpts = max_keep_ckpts |
| | self.save_last = save_last |
| | self.args = kwargs |
| | self.sync_buffer = sync_buffer |
| | self.file_client_args = file_client_args |
| |
|
| | def before_run(self, runner): |
| | if not self.out_dir: |
| | self.out_dir = runner.work_dir |
| |
|
| | self.file_client = FileClient.infer_client(self.file_client_args, |
| | self.out_dir) |
| |
|
| | |
| | |
| | |
| | |
| | if self.out_dir != runner.work_dir: |
| | basename = osp.basename(runner.work_dir.rstrip(osp.sep)) |
| | self.out_dir = self.file_client.join_path(self.out_dir, basename) |
| |
|
| | runner.logger.info((f'Checkpoints will be saved to {self.out_dir} by ' |
| | f'{self.file_client.name}.')) |
| |
|
| | |
| | |
| | if 'create_symlink' in self.args: |
| | if self.args[ |
| | 'create_symlink'] and not self.file_client.allow_symlink: |
| | self.args['create_symlink'] = False |
| | warnings.warn( |
| | ('create_symlink is set as True by the user but is changed' |
| | 'to be False because creating symbolic link is not ' |
| | f'allowed in {self.file_client.name}')) |
| | else: |
| | self.args['create_symlink'] = self.file_client.allow_symlink |
| |
|
| | def after_train_epoch(self, runner): |
| | if not self.by_epoch: |
| | return |
| |
|
| | |
| | |
| | |
| | if self.every_n_epochs( |
| | runner, self.interval) or (self.save_last |
| | and self.is_last_epoch(runner)): |
| | runner.logger.info( |
| | f'Saving checkpoint at {runner.epoch + 1} epochs') |
| | if self.sync_buffer: |
| | allreduce_params(runner.model.buffers()) |
| | self._save_checkpoint(runner) |
| |
|
| | @master_only |
| | def _save_checkpoint(self, runner): |
| | """Save the current checkpoint and delete unwanted checkpoint.""" |
| | 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'] = self.file_client.join_path( |
| | self.out_dir, cur_ckpt_filename) |
| | |
| | if self.max_keep_ckpts > 0: |
| | if self.by_epoch: |
| | name = 'epoch_{}.pth' |
| | current_ckpt = runner.epoch + 1 |
| | else: |
| | name = 'iter_{}.pth' |
| | current_ckpt = runner.iter + 1 |
| | redundant_ckpts = range( |
| | current_ckpt - self.max_keep_ckpts * self.interval, 0, |
| | -self.interval) |
| | filename_tmpl = self.args.get('filename_tmpl', name) |
| | for _step in redundant_ckpts: |
| | ckpt_path = self.file_client.join_path( |
| | self.out_dir, filename_tmpl.format(_step)) |
| | if self.file_client.isfile(ckpt_path): |
| | self.file_client.remove(ckpt_path) |
| | else: |
| | break |
| |
|
| | def after_train_iter(self, runner): |
| | if self.by_epoch: |
| | return |
| |
|
| | |
| | |
| | |
| | if self.every_n_iters( |
| | runner, self.interval) or (self.save_last |
| | and self.is_last_iter(runner)): |
| | runner.logger.info( |
| | f'Saving checkpoint at {runner.iter + 1} iterations') |
| | if self.sync_buffer: |
| | allreduce_params(runner.model.buffers()) |
| | self._save_checkpoint(runner) |
| |
|
