def parge_config(): parser = argparse.ArgumentParser(description='UDA_TP training') parser.add_argument('config', help='train config file path') parser.add_argument('--work-dir', help='the dir to save logs and models', default='') parser.add_argument('--resume-from', help='the checkpoint file to resume from') parser.add_argument('--launcher', type=str, choices=['none', 'pytorch', 'slurm'], default='none', help='job launcher') parser.add_argument('--tcp-port', type=str, default='5017') parser.add_argument('--set', dest='set_cfgs', default=None, nargs=argparse.REMAINDER, help='set extra config keys if needed') args = parser.parse_args() cfg_from_yaml_file(args.config, cfg) cfg.launcher = args.launcher cfg.tcp_port = args.tcp_port if not args.work_dir: args.work_dir = Path(args.config).stem cfg.work_dir = cfg.LOGS_ROOT / args.work_dir mkdir_if_missing(cfg.work_dir) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs, cfg) shutil.copy(args.config, cfg.work_dir / 'config.yaml') return args, cfg
def parge_config(): parser = argparse.ArgumentParser(description="CycleGAN training") parser.add_argument("config", help="train config file path") parser.add_argument("--work-dir", help="the dir to save logs and models", default="") parser.add_argument("--resume-from", help="the checkpoint file to resume from") parser.add_argument( "--launcher", type=str, choices=["none", "pytorch", "slurm"], default="none", help="job launcher", ) parser.add_argument("--tcp-port", type=str, default="5017") parser.add_argument( "--set", dest="set_cfgs", default=None, nargs=argparse.REMAINDER, help="set extra config keys if needed", ) args = parser.parse_args() cfg_from_yaml_file(args.config, cfg) assert len(list(cfg.TRAIN.datasets.keys()))==2, \ "the number of datasets for domain-translation training should be two" cfg.launcher = args.launcher cfg.tcp_port = args.tcp_port if not args.work_dir: args.work_dir = Path(args.config).stem cfg.work_dir = cfg.LOGS_ROOT / args.work_dir mkdir_if_missing(cfg.work_dir) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs, cfg) shutil.copy(args.config, cfg.work_dir / "config.yaml") return args, cfg