Vincentqyw
fix: roma
c74a070
raw
history blame
3 kB
import argparse
arg_lists = []
parser = argparse.ArgumentParser(description="LANet")
def str2bool(v):
return v.lower() in ("true", "1")
def add_argument_group(name):
arg = parser.add_argument_group(name)
arg_lists.append(arg)
return arg
# train data params
traindata_arg = add_argument_group("Traindata Params")
traindata_arg.add_argument("--train_txt", type=str, default="", help="Train set.")
traindata_arg.add_argument(
"--train_root", type=str, default="", help="Where the train images are."
)
traindata_arg.add_argument(
"--batch_size", type=int, default=8, help="# of images in each batch of data"
)
traindata_arg.add_argument(
"--num_workers",
type=int,
default=4,
help="# of subprocesses to use for data loading",
)
traindata_arg.add_argument(
"--pin_memory",
type=str2bool,
default=True,
help="# of subprocesses to use for data loading",
)
traindata_arg.add_argument(
"--shuffle",
type=str2bool,
default=True,
help="Whether to shuffle the train and valid indices",
)
traindata_arg.add_argument("--image_shape", type=tuple, default=(240, 320), help="")
traindata_arg.add_argument(
"--jittering", type=tuple, default=(0.5, 0.5, 0.2, 0.05), help=""
)
# data storage
storage_arg = add_argument_group("Storage")
storage_arg.add_argument("--ckpt_name", type=str, default="PointModel", help="")
# training params
train_arg = add_argument_group("Training Params")
train_arg.add_argument("--start_epoch", type=int, default=0, help="")
train_arg.add_argument("--max_epoch", type=int, default=12, help="")
train_arg.add_argument(
"--init_lr", type=float, default=3e-4, help="Initial learning rate value."
)
train_arg.add_argument(
"--lr_factor", type=float, default=0.5, help="Reduce learning rate value."
)
train_arg.add_argument(
"--momentum", type=float, default=0.9, help="Nesterov momentum value."
)
train_arg.add_argument("--display", type=int, default=50, help="")
# loss function params
loss_arg = add_argument_group("Loss function Params")
loss_arg.add_argument("--score_weight", type=float, default=1.0, help="")
loss_arg.add_argument("--loc_weight", type=float, default=1.0, help="")
loss_arg.add_argument("--desc_weight", type=float, default=4.0, help="")
loss_arg.add_argument("--corres_weight", type=float, default=0.5, help="")
loss_arg.add_argument("--corres_threshold", type=int, default=4.0, help="")
# other params
misc_arg = add_argument_group("Misc.")
misc_arg.add_argument(
"--use_gpu", type=str2bool, default=True, help="Whether to run on the GPU."
)
misc_arg.add_argument("--gpu", type=int, default=0, help="Which GPU to run on.")
misc_arg.add_argument(
"--seed", type=int, default=1001, help="Seed to ensure reproducibility."
)
misc_arg.add_argument(
"--ckpt_dir",
type=str,
default="./checkpoints",
help="Directory in which to save model checkpoints.",
)
def get_config():
config, unparsed = parser.parse_known_args()
return config, unparsed