Spaces:
Sleeping
Sleeping
File size: 2,845 Bytes
761be79 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
# 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)
|