| import argparse |
| import os |
|
|
| import yaml |
|
|
| __all__ = ['get_config', 'print_config'] |
|
|
|
|
| def get_config(args): |
|
|
| config = dict2namespace(setdefault(_get_raw_config(args.config), _get_raw_config("default.yml"))) |
|
|
| if not hasattr(config.sampling, "sigma_dist"): |
| config.sampling.sigma_dist = config.model.sigma_dist |
| if not hasattr(config.biggan, "resolution"): |
| config.biggan.resolution = config.data.image_size |
|
|
| if args.consistent: |
| config.sampling.consistent = args.consistent |
| config.sampling.noise_first = False |
| if args.step_lr: |
| config.sampling.step_lr = args.step_lr |
| if args.nsigma != 0: |
| config.sampling.nsigma = args.nsigma |
| if args.step_lr != 0: |
| config.sampling.step_lr = args.step_lr |
| if args.batch_size != 0: |
| config.sampling.batch_size = args.batch_size |
| config.fast_fid.batch_size = args.batch_size |
| |
| if args.model_types is not None and len(args.model_types)==1 and args.model_types[0] in [0, 6, 23] and config.data.dataset in ['tinyImages', 'CIFAR10']: |
| config.sampling.batch_size = min(200, config.sampling.batch_size) |
| if args.model_types is not None and len(args.model_types) == 1 and args.model_types[0] in [8] and config.data.dataset in ['tinyImages', 'CIFAR10']: |
| config.sampling.batch_size = min(800, config.sampling.batch_size) |
|
|
| if args.ODI_steps == -1: |
| args.ODI_steps = None |
| if args.fid_num_samples != 0: |
| config.fast_fid.num_samples = args.fid_num_samples |
| if args.begin_ckpt != 0: |
| config.fast_fid.begin_ckpt = args.begin_ckpt |
| config.sampling.ckpt_id = args.begin_ckpt |
| if args.end_ckpt != 0: |
| config.fast_fid.end_ckpt = args.begin_ckpt |
| if args.adam: |
| config.optim.beta1 = args.adam_beta[0] |
| config.optim.beta2 = args.adam_beta[1] |
| if args.D_adam: |
| config.optim.adv_beta1 = args.D_adam_beta[0] |
| config.optim.adv_beta2 = args.D_adam_beta[1] |
| if args.D_steps != 0: |
| config.adversarial.D_steps = args.D_steps |
|
|
| return config |
|
|
|
|
| def _get_raw_config(name): |
| here = os.path.dirname(os.path.abspath(__file__)) |
| with open(os.path.join(here, name), 'r') as f: |
| yaml_dict = yaml.load(f, Loader=yaml.FullLoader) |
| return yaml_dict |
|
|
|
|
| def setdefault(config, default): |
| |
| for x in default: |
| v = default.get(x) |
| if isinstance(v, dict) and x in config: |
| setdefault(config.get(x), v) |
| else: |
| config.setdefault(x, v) |
| return config |
|
|
|
|
| def dict2namespace(config): |
| namespace = argparse.Namespace() |
| for key, value in config.items(): |
| if isinstance(value, dict): |
| new_value = dict2namespace(value) |
| else: |
| new_value = value |
| setattr(namespace, key, new_value) |
| return namespace |
|
|
|
|
| def print_config(config): |
| print(">" * 80) |
| print(yaml.dump(config, default_flow_style=False)) |
| print("<" * 80) |
|
|
|
|
