Spaces:
Running on Zero
Running on Zero
| import copy | |
| from mmcv.runner.hooks import CheckpointHook as _CheckpointHook | |
| from mmcv.runner import HOOKS | |
| from mmcv.runner.dist_utils import allreduce_params, get_dist_info | |
| from lakonlab.parallel import FSDPWrapper | |
| class CheckpointHook(_CheckpointHook): | |
| def __init__(self, | |
| interval: int = -1, | |
| must_save_interval: int = 1000000000, | |
| **kwargs): | |
| super().__init__(interval=interval, **kwargs) | |
| self.must_save_interval = must_save_interval | |
| def after_train_epoch(self, runner): | |
| if not self.by_epoch: | |
| return | |
| if self.every_n_epochs(runner, self.interval) \ | |
| or self.every_n_epochs(runner, self.must_save_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, asynchronous=not self.is_last_epoch(runner)) | |
| def after_train_iter(self, runner): | |
| if self.by_epoch: | |
| return | |
| if self.every_n_iters(runner, self.interval) \ | |
| or self.every_n_iters(runner, self.must_save_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, asynchronous=not self.is_last_iter(runner)) | |
| def _save_checkpoint(self, runner, asynchronous=False): | |
| rank, _ = get_dist_info() | |
| # in FSDP we need to call save_checkpoint on all ranks | |
| if rank == 0 or isinstance(runner.model, FSDPWrapper): | |
| if self.max_keep_ckpts > 0: | |
| # remove other checkpoints | |
| cur_progress = copy.deepcopy(runner.epoch) if self.by_epoch else copy.deepcopy(runner.iter) | |
| def after_save_hook(): | |
| if self.by_epoch: | |
| name = 'epoch_{}.pth' | |
| current_ckpt = cur_progress + 1 | |
| else: | |
| name = 'iter_{}.pth' | |
| current_ckpt = cur_progress + 1 | |
| redundant_ckpts = [i for i in range( | |
| current_ckpt - self.max_keep_ckpts * self.interval, 0, | |
| -self.interval) if i % self.must_save_interval != 0] | |
| 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 | |
| else: | |
| after_save_hook = None | |
| runner.save_checkpoint( | |
| self.out_dir, | |
| save_optimizer=self.save_optimizer, | |
| after_save_hook=after_save_hook, | |
| asynchronous=asynchronous, | |
| **self.args) | |
| if rank == 0: | |
| 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) | |