| |
| 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) |
|
|
