90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
from modelscope.utils.config import Config
|
|
|
|
DEFAULT_HOOKS_CONFIG = {
|
|
'train.hooks': [{
|
|
'type': 'CheckpointHook',
|
|
'interval': 1
|
|
}, {
|
|
'type': 'TextLoggerHook',
|
|
'interval': 10
|
|
}, {
|
|
'type': 'IterTimerHook'
|
|
}]
|
|
}
|
|
|
|
_HOOK_KEY_CHAIN_MAP = {
|
|
'TextLoggerHook': 'train.logging',
|
|
'CheckpointHook': 'train.checkpoint.period',
|
|
'BestCkptSaverHook': 'train.checkpoint.best',
|
|
'EvaluationHook': 'evaluation.period',
|
|
}
|
|
|
|
|
|
def merge_cfg(cfg: Config):
|
|
"""Merge the default config into the input cfg.
|
|
|
|
This function will pop the default CheckpointHook when the BestCkptSaverHook exists in the input cfg.
|
|
|
|
Aegs:
|
|
cfg: The input cfg to be merged into.
|
|
"""
|
|
cfg.merge_from_dict(DEFAULT_HOOKS_CONFIG, force=False)
|
|
|
|
|
|
def merge_hooks(cfg: Config) -> List[Dict]:
|
|
hooks = getattr(cfg.train, 'hooks', []).copy()
|
|
for hook_type, key_chain in _HOOK_KEY_CHAIN_MAP.items():
|
|
hook = _key_chain_to_hook(cfg, key_chain, hook_type)
|
|
if hook is not None:
|
|
hooks.append(hook)
|
|
return hooks
|
|
|
|
|
|
def update_cfg(cfg: Config) -> Config:
|
|
if 'hooks' not in cfg.train:
|
|
return cfg
|
|
key_chain_map = {}
|
|
for hook in cfg.train.hooks:
|
|
if not hook:
|
|
continue
|
|
key, value = _hook_split(hook)
|
|
if key not in _HOOK_KEY_CHAIN_MAP:
|
|
continue
|
|
key_chain_map[_HOOK_KEY_CHAIN_MAP[key]] = value
|
|
hook.clear()
|
|
cfg.train.hooks = list(filter(bool, cfg.train.hooks))
|
|
cfg.merge_from_dict(key_chain_map, force=False)
|
|
return cfg
|
|
|
|
|
|
def _key_chain_to_hook(cfg: Config, key_chain: str,
|
|
hook_type: str) -> Optional[Dict]:
|
|
if not _check_basic_hook(cfg, key_chain, hook_type):
|
|
return None
|
|
hook_params: Dict = cfg.safe_get(key_chain)
|
|
hook = {'type': hook_type}
|
|
hook.update(hook_params)
|
|
return hook
|
|
|
|
|
|
def _check_basic_hook(cfg: Config, key_chain: str, hook_type: str) -> bool:
|
|
if cfg.safe_get(key_chain) is None:
|
|
return False
|
|
hooks = list(
|
|
filter(lambda hook: hook['type'] == hook_type,
|
|
getattr(cfg.train, 'hooks', [])))
|
|
assert len(hooks) == 0, f'The key_chain {key_chain} and the traditional hook ' \
|
|
f'cannot exist at the same time, ' \
|
|
f'please delete {hook_type} in the configuration file.'
|
|
return True
|
|
|
|
|
|
def _hook_split(hook: Dict) -> Tuple[str, Dict]:
|
|
hook = hook.copy()
|
|
return hook.pop('type'), hook
|