import copy import logging import os import os.path as osp import pickle import platform import time import warnings from collections import OrderedDict from functools import partial from typing import Callable, Dict, List, Optional, Sequence, Union import torch import torch.nn as nn from lightning.pytorch.loggers import Logger from torch.nn.parallel.distributed import DistributedDataParallel from torch.optim import Optimizer from torch.utils.data import DataLoader import mmengine from mmengine.config import Config, ConfigDict from mmengine.dataset import worker_init_fn from mmengine.device import get_device from mmengine.dist import (broadcast, get_dist_info, get_rank, init_dist, is_distributed, master_only) from mmengine.evaluator import Evaluator from mmengine.fileio import FileClient, join_path from mmengine.hooks import Hook from mmengine.logging import MessageHub, MMLogger, print_log from mmengine.model import (MMDistributedDataParallel, convert_sync_batchnorm, is_model_wrapper, revert_sync_batchnorm) from mmengine.optim import (OptimWrapper, OptimWrapperDict, _ParamScheduler, build_optim_wrapper) from mmengine.registry import (DATA_SAMPLERS, DATASETS, EVALUATOR, FUNCTIONS, HOOKS, LOG_PROCESSORS, LOOPS, MODEL_WRAPPERS, OPTIM_WRAPPERS, PARAM_SCHEDULERS, RUNNERS, VISUALIZERS, DefaultScope) from mmengine.utils import digit_version, get_git_hash, is_seq_of from mmengine.utils.dl_utils import (TORCH_VERSION, collect_env, set_multi_processing) from mmengine.visualization import Visualizer from mmengine.runner.base_loop import BaseLoop from mmengine.runner.checkpoint import (_load_checkpoint, _load_checkpoint_to_model, find_latest_checkpoint, get_state_dict, save_checkpoint, weights_to_cpu) from mmengine.runner.log_processor import LogProcessor from mmengine.runner.loops import EpochBasedTrainLoop, IterBasedTrainLoop, TestLoop, ValLoop from mmengine.runner.priority import Priority, get_priority from mmengine.runner.utils import set_random_seed ConfigType = Union[Dict, Config, ConfigDict] ParamSchedulerType = Union[List[_ParamScheduler], Dict[str, List[_ParamScheduler]]] OptimWrapperType = Union[OptimWrapper, OptimWrapperDict] from mmpl.registry import MODELS, LOGGERS import lightning.pytorch as pl from mmpl.models import build_pler @RUNNERS.register_module() class PLRunner: def __init__( self, trainer_cfg: Dict, model_cfg: Union[pl.LightningModule, Dict], datamodule_cfg: Optional[Dict] = None, cfg: Optional[ConfigType] = None ): self.trainer_cfg = copy.deepcopy(trainer_cfg) self.model_cfg = copy.deepcopy(model_cfg) self.datamodule_cfg = copy.deepcopy(datamodule_cfg) mmengine.mkdir_or_exist(trainer_cfg['default_root_dir']) timestamp = torch.tensor(time.time(), dtype=torch.float64) # broadcast timestamp from 0 process to other processes broadcast(timestamp) self.timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime(timestamp.item())) if cfg is not None: if isinstance(cfg, Config): self.cfg = copy.deepcopy(cfg) elif isinstance(cfg, dict): self.cfg = Config(cfg) else: self.cfg = Config(dict()) compiled_model = trainer_cfg.pop('compiled_model', False) # build logger loggers = self.build_logger( trainer_cfg.get('logger', False), trainer_cfg.get('default_root_dir', f'{self.timestamp}') ) trainer_cfg['logger'] = loggers # build visualizer used for writing log or visualizing all kinds of data self.visualizer = self.build_visualizer( self.cfg.get('visualizer', None), trainer_cfg.get('default_root_dir', f'{self.timestamp}') ) if self.cfg: self.visualizer.add_config(self.cfg) # build callbacks callbacks = self.build_hooks( trainer_cfg.get('callbacks', None), ) trainer_cfg['callbacks'] = callbacks # build strategy strategy = self.build_strategy( trainer_cfg.get('strategy', 'auto'), ) trainer_cfg['strategy'] = strategy self.trainer = pl.Trainer(**trainer_cfg) model_cfg.update({'config_cfg': copy.deepcopy(cfg).to_dict()}) model = self.build_model(model_cfg) if cfg.get('load_from', None) is not None: self.load_checkpoint(model, cfg['load_from']) if compiled_model: # default, reduce-overhead, and max-autotune. self.model = torch.compile(model) else: self.model = model # dump `cfg` to `work_dir` self.dump_config() # # Collect and log environment information. # self._log_env(env_cfg) # log hooks information # self.logger.info(f'Hooks will be executed in the following ' # f'order:\n{self.get_hooks_info()}') def build_visualizer( self, visualizer: Optional[Union[Visualizer, Dict]] = None, default_root_dir = 'tmp' ) -> Visualizer: """Build a global asscessable Visualizer. Args: visualizer (Visualizer or dict, optional): A Visualizer object or a dict to build Visualizer object. If ``visualizer`` is a Visualizer object, just returns itself. If not specified, default config will be used to build Visualizer object. Defaults to None. Returns: Visualizer: A Visualizer object build from ``visualizer``. """ if visualizer is None: visualizer = dict( name=os.path.basename(default_root_dir), vis_backends=[dict(type='LocalVisBackend')], save_dir=default_root_dir+'/visualizer' ) return Visualizer.get_instance(**visualizer) if isinstance(visualizer, Visualizer): return visualizer if isinstance(visualizer, dict): # ensure visualizer containing name key visualizer.setdefault('name', os.path.basename(default_root_dir)) visualizer.setdefault('save_dir', default_root_dir+'/visualizer') return VISUALIZERS.build(visualizer) else: raise TypeError( 'visualizer should be Visualizer object, a dict or None, ' f'but got {visualizer}') def build_hooks(self, hooks: Union[Dict, List[Dict]] = None) -> List[Hook]: """Build hooks from config. Args: hooks_cfg (dict): Config dict of hooks. Returns: list[Hook]: A list of hooks. """ if hooks is not None: if isinstance(hooks, dict): hooks = [hooks] tmp_hooks = [] for hook in hooks: hook = HOOKS.build(hook) tmp_hooks.append(hook) hooks = tmp_hooks return hooks @classmethod def from_cfg(cls, cfg: ConfigType) -> 'Runner': cfg = copy.deepcopy(cfg) runner = cls( trainer_cfg=cfg.get('trainer_cfg'), model_cfg=cfg['model_cfg'], datamodule_cfg=cfg.get('datamodule_cfg'), cfg=cfg ) return runner def build_logger(self, loggers: Union[Dict, List[Dict]] = None, default_root_dir='logger'): if loggers is not None and loggers: if isinstance(loggers, Dict): loggers = [loggers] tmp_loggers = [] for logger in loggers: if logger.get('save_dir', None) is None: logger['save_dir'] = default_root_dir mmengine.mkdir_or_exist(logger['save_dir']) tmp_loggers.append(LOGGERS.build(logger)) loggers = tmp_loggers return loggers def build_strategy(self, strategy='auto'): if isinstance(strategy, str): return strategy elif isinstance(strategy, dict): if strategy.get('type', '') == 'FSDPStrategy': from torch.distributed.fsdp import CPUOffload from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy import functools strategy.update( dict( # cpu_offload=CPUOffload(offload_params=True), auto_wrap_policy=functools.partial( size_based_auto_wrap_policy, min_num_params=int(5e7) ) ) ) strategy = MODEL_WRAPPERS.build(strategy) return strategy return strategy def build_model(self, model: Union[pl.LightningModule, Dict]) -> pl.LightningModule: if isinstance(model, pl.LightningModule): return model elif isinstance(model, dict): model = build_pler(model) return model # type: ignore else: raise TypeError('model should be a nn.Module object or dict, ' f'but got {model}') def _init_model_weights(self) -> None: """Initialize the model weights if the model has :meth:`init_weights`""" if hasattr(self.model, 'module'): model = self.model.module else: model = self.model if hasattr(model, 'init_weights'): model.init_weights() # sync params and buffers for name, params in model.state_dict().items(): broadcast(params) def get_hooks_info(self) -> str: # Get hooks info in each stage stage_hook_map: Dict[str, list] = {stage: [] for stage in Hook.stages} for hook in self.hooks: try: priority = Priority(hook.priority).name # type: ignore except ValueError: priority = hook.priority # type: ignore classname = hook.__class__.__name__ hook_info = f'({priority:<12}) {classname:<35}' for trigger_stage in hook.get_triggered_stages(): stage_hook_map[trigger_stage].append(hook_info) stage_hook_infos = [] for stage in Hook.stages: hook_infos = stage_hook_map[stage] if len(hook_infos) > 0: info = f'{stage}:\n' info += '\n'.join(hook_infos) info += '\n -------------------- ' stage_hook_infos.append(info) return '\n'.join(stage_hook_infos) def load_or_resume(self) -> None: """load or resume checkpoint.""" if self._has_loaded: return None # decide to load from checkpoint or resume from checkpoint resume_from = None if self._resume and self._load_from is None: # auto resume from the latest checkpoint resume_from = find_latest_checkpoint(self.work_dir) self.logger.info( f'Auto resumed from the latest checkpoint {resume_from}.') elif self._resume and self._load_from is not None: # resume from the specified checkpoint resume_from = self._load_from if resume_from is not None: self.resume(resume_from) self._has_loaded = True elif self._load_from is not None: self.load_checkpoint(self._load_from) self._has_loaded = True @staticmethod def build_datamodule(datamodule_cfg: Union[pl.LightningDataModule, Dict]): if isinstance(datamodule_cfg, pl.LightningDataModule): return datamodule_cfg datamodule_cfg = copy.deepcopy(datamodule_cfg) # build datamodule datamodule = DATASETS.build(datamodule_cfg) return datamodule def run(self, status, *args, **kwargs): assert status in ['fit', 'test', 'predict', 'validate'] trainer_func = self.trainer.__getattribute__(status) self.datamodule = self.build_datamodule(self.datamodule_cfg) return trainer_func(model=self.model, datamodule=self.datamodule, *args, **kwargs) # # if is_model_wrapper(self.model): # ori_model = self.model.module # else: # ori_model = self.model # assert hasattr(ori_model, 'train_step'), ( # 'If you want to train your model, please make sure your model ' # 'has implemented `train_step`.') # # if self._val_loop is not None: # assert hasattr(ori_model, 'val_step'), ( # 'If you want to validate your model, please make sure your ' # 'model has implemented `val_step`.') # # if self._train_loop is None: # raise RuntimeError( # '`self._train_loop` should not be None when calling train ' # 'method. Please provide `train_dataloader`, `train_cfg`, ' # '`optimizer` and `param_scheduler` arguments when ' # 'initializing runner.') # # self._train_loop = self.build_train_loop( # self._train_loop) # type: ignore # # # `build_optimizer` should be called before `build_param_scheduler` # # because the latter depends on the former # self.optim_wrapper = self.build_optim_wrapper(self.optim_wrapper) # # Automatically scaling lr by linear scaling rule # self.scale_lr(self.optim_wrapper, self.auto_scale_lr) # # if self.param_schedulers is not None: # self.param_schedulers = self.build_param_scheduler( # type: ignore # self.param_schedulers) # type: ignore # # if self._val_loop is not None: # self._val_loop = self.build_val_loop( # self._val_loop) # type: ignore # # TODO: add a contextmanager to avoid calling `before_run` many times # self.call_hook('before_run') # # # initialize the model weights # self._init_model_weights() # # make sure checkpoint-related hooks are triggered after `before_run` # self.load_or_resume() # # # Initiate inner count of `optim_wrapper`. # self.optim_wrapper.initialize_count_status( # self.model, # self._train_loop.iter, # type: ignore # self._train_loop.max_iters) # type: ignore # # # Maybe compile the model according to options in self.cfg.compile # # This must be called **AFTER** model has been wrapped. # self._maybe_compile('train_step') # # model = self.train_loop.run() # type: ignore # self.call_hook('after_run') # return model def register_hook( self, hook: Union[Hook, Dict], priority: Optional[Union[str, int, Priority]] = None) -> None: """Register a hook into the hook list. The hook will be inserted into a priority queue, with the specified priority (See :class:`Priority` for details of priorities). For hooks with the same priority, they will be triggered in the same order as they are registered. Priority of hook will be decided with the following priority: - ``priority`` argument. If ``priority`` is given, it will be priority of hook. - If ``hook`` argument is a dict and ``priority`` in it, the priority will be the value of ``hook['priority']``. - If ``hook`` argument is a dict but ``priority`` not in it or ``hook`` is an instance of ``hook``, the priority will be ``hook.priority``. Args: hook (:obj:`Hook` or dict): The hook to be registered. priority (int or str or :obj:`Priority`, optional): Hook priority. Lower value means higher priority. """ if not isinstance(hook, (Hook, dict)): raise TypeError( f'hook should be an instance of Hook or dict, but got {hook}') _priority = None if isinstance(hook, dict): if 'priority' in hook: _priority = hook.pop('priority') hook_obj = HOOKS.build(hook) else: hook_obj = hook if priority is not None: hook_obj.priority = priority elif _priority is not None: hook_obj.priority = _priority inserted = False for i in range(len(self._hooks) - 1, -1, -1): if get_priority(hook_obj.priority) >= get_priority( self._hooks[i].priority): self._hooks.insert(i + 1, hook_obj) inserted = True break if not inserted: self._hooks.insert(0, hook_obj) def register_default_hooks( self, hooks: Optional[Dict[str, Union[Hook, Dict]]] = None) -> None: """Register default hooks into hook list. ``hooks`` will be registered into runner to execute some default actions like updating model parameters or saving checkpoints. Default hooks and their priorities: +----------------------+-------------------------+ | Hooks | Priority | +======================+=========================+ | RuntimeInfoHook | VERY_HIGH (10) | +----------------------+-------------------------+ | IterTimerHook | NORMAL (50) | +----------------------+-------------------------+ | DistSamplerSeedHook | NORMAL (50) | +----------------------+-------------------------+ | LoggerHook | BELOW_NORMAL (60) | +----------------------+-------------------------+ | ParamSchedulerHook | LOW (70) | +----------------------+-------------------------+ | CheckpointHook | VERY_LOW (90) | +----------------------+-------------------------+ If ``hooks`` is None, above hooks will be registered by default:: default_hooks = dict( runtime_info=dict(type='RuntimeInfoHook'), timer=dict(type='IterTimerHook'), sampler_seed=dict(type='DistSamplerSeedHook'), logger=dict(type='LoggerHook'), param_scheduler=dict(type='ParamSchedulerHook'), checkpoint=dict(type='CheckpointHook', interval=1), ) If not None, ``hooks`` will be merged into ``default_hooks``. If there are None value in default_hooks, the corresponding item will be popped from ``default_hooks``:: hooks = dict(timer=None) The final registered default hooks will be :obj:`RuntimeInfoHook`, :obj:`DistSamplerSeedHook`, :obj:`LoggerHook`, :obj:`ParamSchedulerHook` and :obj:`CheckpointHook`. Args: hooks (dict[str, Hook or dict], optional): Default hooks or configs to be registered. """ default_hooks: dict = dict( runtime_info=dict(type='RuntimeInfoHook'), timer=dict(type='IterTimerHook'), sampler_seed=dict(type='DistSamplerSeedHook'), logger=dict(type='LoggerHook'), param_scheduler=dict(type='ParamSchedulerHook'), checkpoint=dict(type='CheckpointHook', interval=1), ) if hooks is not None: for name, hook in hooks.items(): if name in default_hooks and hook is None: # remove hook from _default_hooks default_hooks.pop(name) else: assert hook is not None default_hooks[name] = hook for hook in default_hooks.values(): self.register_hook(hook) def register_custom_hooks(self, hooks: List[Union[Hook, Dict]]) -> None: """Register custom hooks into hook list. Args: hooks (list[Hook | dict]): List of hooks or configs to be registered. """ for hook in hooks: self.register_hook(hook) def register_hooks( self, default_hooks: Optional[Dict[str, Union[Hook, Dict]]] = None, custom_hooks: Optional[List[Union[Hook, Dict]]] = None) -> None: """Register default hooks and custom hooks into hook list. Args: default_hooks (dict[str, dict] or dict[str, Hook], optional): Hooks to execute default actions like updating model parameters and saving checkpoints. Defaults to None. custom_hooks (list[dict] or list[Hook], optional): Hooks to execute custom actions like visualizing images processed by pipeline. Defaults to None. """ self.register_default_hooks(default_hooks) if custom_hooks is not None: self.register_custom_hooks(custom_hooks) def resume(self, filename: str, resume_optimizer: bool = True, resume_param_scheduler: bool = True, map_location: Union[str, Callable] = 'default') -> None: """Resume model from checkpoint. Args: filename (str): Accept local filepath, URL, ``torchvision://xxx``, ``open-mmlab://xxx``. resume_optimizer (bool): Whether to resume optimizer state. Defaults to True. resume_param_scheduler (bool): Whether to resume param scheduler state. Defaults to True. map_location (str or callable):A string or a callable function to specifying how to remap storage locations. Defaults to 'default'. """ if map_location == 'default': device = get_device() checkpoint = self.load_checkpoint(filename, map_location=device) else: checkpoint = self.load_checkpoint( filename, map_location=map_location) self.train_loop._epoch = checkpoint['meta']['epoch'] self.train_loop._iter = checkpoint['meta']['iter'] # check whether the number of GPU used for current experiment # is consistent with resuming from checkpoint if 'config' in checkpoint['meta']: config = mmengine.Config.fromstring( checkpoint['meta']['config'], file_format='.py') previous_gpu_ids = config.get('gpu_ids', None) if (previous_gpu_ids is not None and len(previous_gpu_ids) > 0 and len(previous_gpu_ids) != self._world_size): # TODO, should we modify the iteration? self.logger.info( 'Number of GPU used for current experiment is not ' 'consistent with resuming from checkpoint') if (self.auto_scale_lr is None or not self.auto_scale_lr.get('enable', False)): raise RuntimeError( 'Cannot automatically rescale lr in resuming. Please ' 'make sure the number of GPU is consistent with the ' 'previous training state resuming from the checkpoint ' 'or set `enable` in `auto_scale_lr to False.') # resume random seed resumed_seed = checkpoint['meta'].get('seed', None) current_seed = self._randomness_cfg.get('seed') if resumed_seed is not None and resumed_seed != current_seed: if current_seed is not None: print_log( f'The value of random seed in the ' f'checkpoint "{resumed_seed}" is ' f'different from the value in ' f'`randomness` config "{current_seed}"', logger='current', level=logging.WARNING) self._randomness_cfg.update(seed=resumed_seed) self.set_randomness(**self._randomness_cfg) resumed_dataset_meta = checkpoint['meta'].get('dataset_meta', None) dataset_meta = getattr(self.train_dataloader.dataset, 'metainfo', None) # `resumed_dataset_meta` and `dataset_meta` could be object like # np.ndarray, which cannot be directly judged as equal or not, # therefore we just compared their dumped results. if pickle.dumps(resumed_dataset_meta) != pickle.dumps(dataset_meta): print_log( 'The dataset metainfo from the resumed checkpoint is ' 'different from the current training dataset, please ' 'check the correctness of the checkpoint or the training ' 'dataset.', logger='current', level=logging.WARNING) self.message_hub.load_state_dict(checkpoint['message_hub']) # resume optimizer if 'optimizer' in checkpoint and resume_optimizer: self.optim_wrapper = self.build_optim_wrapper(self.optim_wrapper) self.optim_wrapper.load_state_dict( # type: ignore checkpoint['optimizer']) # resume param scheduler if resume_param_scheduler and self.param_schedulers is None: print_log( '`resume_param_scheduler` is True but `self.param_schedulers` ' 'is None, so skip resuming parameter schedulers', logger='current', level=logging.WARNING) resume_param_scheduler = False if 'param_schedulers' in checkpoint and resume_param_scheduler: self.param_schedulers = self.build_param_scheduler( # type: ignore self.param_schedulers) # type: ignore if isinstance(self.param_schedulers, dict): for name, schedulers in self.param_schedulers.items(): for scheduler, ckpt_scheduler in zip( schedulers, checkpoint['param_schedulers'][name]): scheduler.load_state_dict(ckpt_scheduler) else: for scheduler, ckpt_scheduler in zip( self.param_schedulers, # type: ignore checkpoint['param_schedulers']): scheduler.load_state_dict(ckpt_scheduler) self._has_loaded = True self.logger.info(f'resumed epoch: {self.epoch}, iter: {self.iter}') # def load_checkpoint(self, # filename: str, # model, # map_location: Union[str, Callable] = 'cpu', # strict: bool = False, # revise_keys: list = [(r'^module.', '')]): # """Load checkpoint from given ``filename``. # # Args: # filename (str): Accept local filepath, URL, ``torchvision://xxx``, # ``open-mmlab://xxx``. # map_location (str or callable): A string or a callable function to # specifying how to remap storage locations. # Defaults to 'cpu'. # strict (bool): strict (bool): Whether to allow different params for # the model and checkpoint. # revise_keys (list): A list of customized keywords to modify the # state_dict in checkpoint. Each item is a (pattern, replacement) # pair of the regular expression operations. Defaults to strip # the prefix 'module.' by [(r'^module\\.', '')]. # """ # checkpoint = _load_checkpoint(filename, map_location=map_location) # # if is_model_wrapper(model): # model = model.module # else: # model = model # # checkpoint = _load_checkpoint_to_model( # model, checkpoint, strict, revise_keys=revise_keys) # # print(f'Load checkpoint from {filename}') # # return checkpoint def load_checkpoint(self, model, file): if isinstance(file, str): file_path = file state_dict = torch.load(file_path, map_location='cpu')['state_dict'] elif isinstance(file, dict): file_path = file['file_path'] state_dict = torch.load(file_path, map_location='cpu')['state_dict'] for delete_key in file['delete_keys']: del state_dict[delete_key] else: raise TypeError('file must be str or dict') missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) print('load from:', file_path) print('load model missing_keys:', missing_keys) print('load model unexpected_keys:', unexpected_keys) @master_only def save_checkpoint( self, out_dir: str, filename: str, file_client_args: Optional[dict] = None, save_optimizer: bool = True, save_param_scheduler: bool = True, meta: dict = None, by_epoch: bool = True, backend_args: Optional[dict] = None, ): """Save checkpoints. ``CheckpointHook`` invokes this method to save checkpoints periodically. Args: out_dir (str): The directory that checkpoints are saved. filename (str): The checkpoint filename. file_client_args (dict, optional): Arguments to instantiate a FileClient. See :class:`mmengine.fileio.FileClient` for details. Defaults to None. It will be deprecated in future. Please use `backend_args` instead. save_optimizer (bool): Whether to save the optimizer to the checkpoint. Defaults to True. save_param_scheduler (bool): Whether to save the param_scheduler to the checkpoint. Defaults to True. meta (dict, optional): The meta information to be saved in the checkpoint. Defaults to None. by_epoch (bool): Whether the scheduled momentum is updated by epochs. Defaults to True. backend_args (dict, optional): Arguments to instantiate the prefix of uri corresponding backend. Defaults to None. New in v0.2.0. """ if meta is None: meta = {} elif not isinstance(meta, dict): raise TypeError( f'meta should be a dict or None, but got {type(meta)}') if by_epoch: # self.epoch increments 1 after # `self.call_hook('after_train_epoch)` but `save_checkpoint` is # called by `after_train_epoch`` method of `CheckpointHook` so # `epoch` should be `self.epoch + 1` meta.update(epoch=self.epoch + 1, iter=self.iter) else: meta.update(epoch=self.epoch, iter=self.iter + 1) if file_client_args is not None: warnings.warn( '"file_client_args" will be deprecated in future. ' 'Please use "backend_args" instead', DeprecationWarning) if backend_args is not None: raise ValueError( '"file_client_args" and "backend_args" cannot be set at ' 'the same time.') file_client = FileClient.infer_client(file_client_args, out_dir) filepath = file_client.join_path(out_dir, filename) else: filepath = join_path( # type: ignore out_dir, filename, backend_args=backend_args) meta.update( cfg=self.cfg.pretty_text, seed=self.seed, experiment_name=self.experiment_name, time=time.strftime('%Y%m%d_%H%M%S', time.localtime()), mmengine_version=mmengine.__version__ + get_git_hash()) if hasattr(self.train_dataloader.dataset, 'metainfo'): meta.update(dataset_meta=self.train_dataloader.dataset.metainfo) if is_model_wrapper(self.model): model = self.model.module else: model = self.model checkpoint = { 'meta': meta, 'state_dict': weights_to_cpu(get_state_dict(model)), 'message_hub': self.message_hub.state_dict() } # save optimizer state dict to checkpoint if save_optimizer: if isinstance(self.optim_wrapper, OptimWrapper): checkpoint['optimizer'] = self.optim_wrapper.state_dict() else: raise TypeError( 'self.optim_wrapper should be an `OptimWrapper` ' 'or `OptimWrapperDict` instance, but got ' f'{self.optim_wrapper}') # save param scheduler state dict if save_param_scheduler and self.param_schedulers is None: print_log( '`save_param_scheduler` is True but `self.param_schedulers` ' 'is None, so skip saving parameter schedulers', logger='current', level=logging.WARNING) save_param_scheduler = False if save_param_scheduler: if isinstance(self.param_schedulers, dict): checkpoint['param_schedulers'] = dict() for name, schedulers in self.param_schedulers.items(): checkpoint['param_schedulers'][name] = [] for scheduler in schedulers: state_dict = scheduler.state_dict() checkpoint['param_schedulers'][name].append(state_dict) else: checkpoint['param_schedulers'] = [] for scheduler in self.param_schedulers: # type: ignore state_dict = scheduler.state_dict() # type: ignore checkpoint['param_schedulers'].append(state_dict) self.call_hook('before_save_checkpoint', checkpoint=checkpoint) save_checkpoint(checkpoint, filepath) @master_only def dump_config(self) -> None: version = '' if len(self.trainer.loggers) > 0: version = self.trainer.loggers[0].version version = version if isinstance(version, str) else f"version_{version}" if version == '': # if no loggers, use default_root_dir version = 'version' """Dump config to `work_dir`.""" if self.cfg.filename is not None: filename = osp.basename(self.cfg.filename) else: filename = f'{self.timestamp}.py' path = f'{self.trainer.default_root_dir}/{version}_{filename}' self.cfg.dump(path) def _check_scheduler_cfg( self, param_scheduler: Optional[Union[dict, list, _ParamScheduler]]) -> None: """Parse `param_scheduler` to a list of parameter schedulers, or a `dict` of which each value is a list of parameter schedulers. If only one optimizer is used, the parsed config should be a list of parameter scheduler configs or instances. If multiple optimizers are used, the parsed config should be `dict`. Its key should be consistent with the optimizer `dict` and its value should be a list of parameter scheduler configs or instances. See :meth:`build_param_scheduler` for more details. Examples: >>> # valid scheduler: >>> # empty scheduler >>> scheduler = None >>> # Single scheduler >>> scheduler = dict(type='MultiStepLR', milestones=[1, 2]) >>> # Single list schedulers >>> scheduler = [dict(type='MultiStepLR', milestones=[1, 2]), >>> dict(type='MultiStepLR', milestones=[2, 3])] >>> # `dict` of schedulers >>> scheduler = dict(linear1=dict(type='MultiStepLR', milestones=[1, 2]), >>> linear2=dict(type='MultiStepLR', milestones=[1, 2])) >>> # `dict` of `list` of schedulers >>> scheduler = dict(linear1=[dict(type='MultiStepLR', milestones=[1, 2])], >>> linear2=[dict(type='MultiStepLR', milestones=[1, 2])]) >>> # Single built scheduler >>> from mmengine.optim import MultiStepLR >>> scheduler = MultiStepLR(milestones=[1, 2], optimizer=optimizer) >>> # Single built list schedulers >>> scheduler = [MultiStepLR(milestones=[1, 2], optimizer=optimizer)] >>> # dict of built scheduler >>> scheduler = dict(linear1=MultiStepLR(milestones=[1, 2], optimizer=optimizer), >>> linear2=MultiStepLR(milestones=[1, 2], optimizer=optimizer)) >>> # dict of built list schedulers >>> scheduler = dict(linear1=[MultiStepLR(milestones=[1, 2], optimizer=optimizer)], >>> linear2=[MultiStepLR(milestones=[1, 2], optimizer=optimizer)]) Args: param_scheduler (dict or list): The original parameter scheduler. """ # noqa: E501 param_schedulers: Union[dict, list, _ParamScheduler] if param_scheduler is None: return if isinstance(param_scheduler, _ParamScheduler): return if is_seq_of(param_scheduler, _ParamScheduler): return if is_seq_of(param_scheduler, dict): for _param_scheduler in param_scheduler: assert 'type' in _param_scheduler, ( 'Each parameter scheduler should contain the key type, ' f'but got {_param_scheduler}') elif isinstance(param_scheduler, dict): if 'type' not in param_scheduler: for key, _param_scheduler in param_scheduler.items(): assert isinstance( _param_scheduler, (dict, tuple, list, _ParamScheduler)), ( 'Each value of `param_scheduler` should be a ' f'dict or a list, but got {_param_scheduler} with ' f'type {type(_ParamScheduler)}') else: raise TypeError( '`param_scheduler` should be a `_ParamScheduler`, `dict`, ' f'list or a tuple, but got {type(param_scheduler)}. If ' '`param_scheduler` is a list of dict, it means a list of ' 'scheduler configs for single optimizer. If it is a dict and ' 'contains key `type`, it means a scheduler config for a ' 'single optimizer. If it does not contain key `type`, it ' 'means multiple lists of schedulers for multiple optimizers.') def _log_env(self, env_cfg: dict) -> None: """Logging environment information of the current task. Args: env_cfg (dict): The environment config of the runner. """ # Collect and log environment information. env = collect_env() runtime_env = OrderedDict() runtime_env.update(env_cfg) runtime_env.update(self._randomness_cfg) runtime_env['Distributed launcher'] = self._launcher runtime_env['Distributed training'] = self._distributed runtime_env['GPU number'] = self._world_size env_info = '\n ' + '\n '.join(f'{k}: {v}' for k, v in env.items()) runtime_env_info = '\n ' + '\n '.join( f'{k}: {v}' for k, v in runtime_env.items()) dash_line = '-' * 60 self.logger.info('\n' + dash_line + '\nSystem environment:' + env_info + '\n' '\nRuntime environment:' + runtime_env_info + '\n' + dash_line + '\n') self.logger.info(f'Config:\n{self.cfg.pretty_text}')