Function train_args

src/utils/parser_util.py:102–145

()

def train_args():
    parser = argparse.ArgumentParser()
    add_base_options(parser)
    add_encoding_training_options(parser)
    add_diffusion_training_options(parser)
    args = parser.parse_args()

    # check existence
    if os.path.exists(args.tag):
        response = input(f'Folder "{args.tag}" already exists, continue? (y/n) ')
        if response != 'y':
            exit()

    os.makedirs(args.tag, exist_ok=True)
    enc_log_dir = encoding_log_dir(args.tag)
    diff_log_dir = diffusion_log_dir(args.tag)

    # encoding part
    if args.enc_log is not None: # use existing encoding, load args
        load_and_overwrite_args(args, os.path.join(args.enc_log, "args.json"))
        if not os.path.exists(enc_log_dir):
            os.symlink(os.path.abspath(args.enc_log), enc_log_dir)
    else:
        os.makedirs(enc_log_dir, exist_ok=True)
        save_path = os.path.join(enc_log_dir, "args.json")
        with open(save_path, "w") as f:
            json.dump(get_args_by_group(parser, args, "encoding"), f, indent=4)
    
    # diffusion part
    args.in_channels = args.fdim_geo if args.data_type == "sdf" else args.fdim_geo + args.fdim_tex
    args.out_channels = args.fdim_geo if args.data_type == "sdf" else args.fdim_geo + args.fdim_tex
    os.makedirs(diff_log_dir, exist_ok=True)
    save_path = os.path.join(diff_log_dir, "args.json")
    with open(save_path, "w") as f:
        json.dump(get_args_by_group(parser, args, "diffusion"), f, indent=4)
    
    # assert args.in_channels == args.out_channels == args.fdim_geo + args.fdim_tex

    # print all args
    print("----- Training args -----")
    for k, v in args.__dict__.items():
        print("{0:20}".format(k), v)

    return args


def sample_args():