Spaces:
Sleeping
Sleeping
| # Omegaconf utils. | |
| # Mostly from https://github.com/bennyguo/instant-nsr-pl/blob/main/utils/misc.py | |
| # Contributer(s): Neil Z. Shao | |
| # All rights reserved. Prometheus 2023. | |
| import os | |
| from omegaconf import OmegaConf | |
| from packaging import version | |
| from pathlib import Path | |
| # ============ Register OmegaConf Recolvers ============= # | |
| OmegaConf.register_new_resolver('calc_exp_lr_decay_rate', lambda factor, n: factor**(1./n)) | |
| OmegaConf.register_new_resolver('add', lambda a, b: a + b) | |
| OmegaConf.register_new_resolver('sub', lambda a, b: a - b) | |
| OmegaConf.register_new_resolver('mul', lambda a, b: a * b) | |
| OmegaConf.register_new_resolver('div', lambda a, b: a / b) | |
| OmegaConf.register_new_resolver('idiv', lambda a, b: a // b) | |
| OmegaConf.register_new_resolver('basename', lambda p: os.path.basename(p)) | |
| # ======================================================= # | |
| OmegaConf.register_new_resolver('int', lambda i: int(i)) | |
| OmegaConf.register_new_resolver('if_else', lambda cond, a, b: a if cond else b) | |
| OmegaConf.register_new_resolver('int2str', lambda x, fmt: f'{x:{fmt}}') | |
| OmegaConf.register_new_resolver('not', lambda a: not a) | |
| # ======================================================= # | |
| def prompt(question): | |
| inp = input(f"{question} (y/n)").lower().strip() | |
| if inp and inp == 'y': | |
| return True | |
| if inp and inp == 'n': | |
| return False | |
| return prompt(question) | |
| def load_config(*yaml_files, cli_args=[]): | |
| yaml_confs = [OmegaConf.load(f) for f in yaml_files] | |
| cli_conf = OmegaConf.from_cli(cli_args) | |
| conf = OmegaConf.merge(*yaml_confs, cli_conf) | |
| OmegaConf.resolve(conf) | |
| # conf.yaml_confs = yaml_confs | |
| conf.yaml_dir = str(Path(yaml_files[0]).parent).replace('\\', '/') | |
| # conf.yaml_dirs = [str(Path(fn).parent).replace('\\', '/') for fn in yaml_files] | |
| return conf | |
| def load_from_config(yaml_files, dat_dir=None, cli_args=[]): | |
| if dat_dir is not None: | |
| for i in range(len(yaml_files)): | |
| # if not os.path.isabs(yaml_files[i]): | |
| # yaml_files[i] = os.path.join(dat_dir, yaml_files[i]) | |
| yaml_files[i] = yaml_files[i].replace('{dat_dir}', dat_dir) | |
| config = load_config(*yaml_files, cli_args=cli_args) | |
| return config | |
| def config_to_primitive(config, resolve=True): | |
| return OmegaConf.to_container(config, resolve=resolve) | |
| def dump_config(path, config): | |
| with open(path, 'w') as fp: | |
| OmegaConf.save(config=config, f=fp) | |
| def get_rank(): | |
| # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing, | |
| # therefore LOCAL_RANK needs to be checked first | |
| rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK") | |
| for key in rank_keys: | |
| rank = os.environ.get(key) | |
| if rank is not None: | |
| return int(rank) | |
| return 0 | |
| def parse_version(ver): | |
| return version.parse(ver) | |