wehub-resource-sync
82 lines
3.1 KiB
Python
82 lines
3.1 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
import os
|
|
import omegaconf
|
|
from omegaconf import OmegaConf
|
|
|
|
|
|
def load_config(args=None, config_file=None, overwrite_fairseq=False):
|
|
"""TODO (huxu): move fairseq overwrite to another function."""
|
|
if args is not None:
|
|
config_file = args.taskconfig
|
|
config = recursive_config(config_file)
|
|
|
|
if config.dataset.subsampling is not None:
|
|
batch_size = config.fairseq.dataset.batch_size // config.dataset.subsampling
|
|
print(
|
|
"adjusting batch_size to {} due to subsampling {}.".format(
|
|
batch_size, config.dataset.subsampling
|
|
)
|
|
)
|
|
config.fairseq.dataset.batch_size = batch_size
|
|
|
|
is_test = config.dataset.split is not None and config.dataset.split == "test"
|
|
if not is_test:
|
|
if (
|
|
config.fairseq.checkpoint is None
|
|
or config.fairseq.checkpoint.save_dir is None
|
|
):
|
|
raise ValueError("fairseq save_dir or save_path must be specified.")
|
|
|
|
save_dir = config.fairseq.checkpoint.save_dir
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
if config.fairseq.common.tensorboard_logdir is not None:
|
|
tb_run_dir = suffix_rundir(
|
|
save_dir, config.fairseq.common.tensorboard_logdir
|
|
)
|
|
config.fairseq.common.tensorboard_logdir = tb_run_dir
|
|
print(
|
|
"update tensorboard_logdir as", config.fairseq.common.tensorboard_logdir
|
|
)
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
OmegaConf.save(config=config, f=os.path.join(save_dir, "config.yaml"))
|
|
|
|
if overwrite_fairseq and config.fairseq is not None and args is not None:
|
|
# flatten fields.
|
|
for group in config.fairseq:
|
|
for field in config.fairseq[group]:
|
|
print("overwrite args." + field, "as", config.fairseq[group][field])
|
|
setattr(args, field, config.fairseq[group][field])
|
|
return config
|
|
|
|
|
|
def recursive_config(config_path):
|
|
"""allows for stacking of configs in any depth."""
|
|
config = OmegaConf.load(config_path)
|
|
if config.includes is not None:
|
|
includes = config.includes
|
|
config.pop("includes")
|
|
base_config = recursive_config(includes)
|
|
config = OmegaConf.merge(base_config, config)
|
|
return config
|
|
|
|
|
|
def suffix_rundir(save_dir, run_dir):
|
|
max_id = -1
|
|
for search_dir in os.listdir(save_dir):
|
|
if search_dir.startswith(run_dir):
|
|
splits = search_dir.split("_")
|
|
cur_id = int(splits[1]) if len(splits) > 1 else 0
|
|
max_id = max(max_id, cur_id)
|
|
return os.path.join(save_dir, run_dir + "_" + str(max_id + 1))
|
|
|
|
|
|
def overwrite_dir(config, replace, basedir):
|
|
for key in config:
|
|
if isinstance(config[key], str) and config[key].startswith(basedir):
|
|
config[key] = config[key].replace(basedir, replace)
|
|
if isinstance(config[key], omegaconf.dictconfig.DictConfig):
|
|
overwrite_dir(config[key], replace, basedir)
|