|
|
|
import os.path as osp |
|
import platform |
|
import shutil |
|
|
|
import torch |
|
from torch.optim import Optimizer |
|
|
|
import mmcv |
|
from mmcv.runner import RUNNERS, EpochBasedRunner |
|
from .checkpoint import save_checkpoint |
|
|
|
try: |
|
import apex |
|
except: |
|
print('apex is not installed') |
|
|
|
|
|
@RUNNERS.register_module() |
|
class EpochBasedRunnerAmp(EpochBasedRunner): |
|
"""Epoch-based Runner with AMP support. |
|
|
|
This runner train models epoch by epoch. |
|
""" |
|
|
|
def save_checkpoint(self, |
|
out_dir, |
|
filename_tmpl='epoch_{}.pth', |
|
save_optimizer=True, |
|
meta=None, |
|
create_symlink=True): |
|
"""Save the checkpoint. |
|
|
|
Args: |
|
out_dir (str): The directory that checkpoints are saved. |
|
filename_tmpl (str, optional): The checkpoint filename template, |
|
which contains a placeholder for the epoch number. |
|
Defaults to 'epoch_{}.pth'. |
|
save_optimizer (bool, optional): Whether to save the optimizer to |
|
the checkpoint. Defaults to True. |
|
meta (dict, optional): The meta information to be saved in the |
|
checkpoint. Defaults to None. |
|
create_symlink (bool, optional): Whether to create a symlink |
|
"latest.pth" to point to the latest checkpoint. |
|
Defaults to True. |
|
""" |
|
if meta is None: |
|
meta = dict(epoch=self.epoch + 1, iter=self.iter) |
|
elif isinstance(meta, dict): |
|
meta.update(epoch=self.epoch + 1, iter=self.iter) |
|
else: |
|
raise TypeError( |
|
f'meta should be a dict or None, but got {type(meta)}') |
|
if self.meta is not None: |
|
meta.update(self.meta) |
|
|
|
filename = filename_tmpl.format(self.epoch + 1) |
|
filepath = osp.join(out_dir, filename) |
|
optimizer = self.optimizer if save_optimizer else None |
|
save_checkpoint(self.model, filepath, optimizer=optimizer, meta=meta) |
|
|
|
|
|
if create_symlink: |
|
dst_file = osp.join(out_dir, 'latest.pth') |
|
if platform.system() != 'Windows': |
|
mmcv.symlink(filename, dst_file) |
|
else: |
|
shutil.copy(filepath, dst_file) |
|
|
|
def resume(self, |
|
checkpoint, |
|
resume_optimizer=True, |
|
map_location='default'): |
|
if map_location == 'default': |
|
if torch.cuda.is_available(): |
|
device_id = torch.cuda.current_device() |
|
checkpoint = self.load_checkpoint( |
|
checkpoint, |
|
map_location=lambda storage, loc: storage.cuda(device_id)) |
|
else: |
|
checkpoint = self.load_checkpoint(checkpoint) |
|
else: |
|
checkpoint = self.load_checkpoint( |
|
checkpoint, map_location=map_location) |
|
|
|
self._epoch = checkpoint['meta']['epoch'] |
|
self._iter = checkpoint['meta']['iter'] |
|
if 'optimizer' in checkpoint and resume_optimizer: |
|
if isinstance(self.optimizer, Optimizer): |
|
self.optimizer.load_state_dict(checkpoint['optimizer']) |
|
elif isinstance(self.optimizer, dict): |
|
for k in self.optimizer.keys(): |
|
self.optimizer[k].load_state_dict( |
|
checkpoint['optimizer'][k]) |
|
else: |
|
raise TypeError( |
|
'Optimizer should be dict or torch.optim.Optimizer ' |
|
f'but got {type(self.optimizer)}') |
|
|
|
if 'amp' in checkpoint: |
|
apex.amp.load_state_dict(checkpoint['amp']) |
|
self.logger.info('load amp state dict') |
|
|
|
self.logger.info('resumed epoch %d, iter %d', self.epoch, self.iter) |
|
|
