49 lines
2.2 KiB
Python
49 lines
2.2 KiB
Python
import os, sys, time, random, argparse
|
|
from .share_args import add_shared_args
|
|
|
|
|
|
def obtain_search_single_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Train a classification model on typical image classification datasets.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
)
|
|
parser.add_argument("--resume", type=str, help="Resume path.")
|
|
parser.add_argument(
|
|
"--model_config", type=str, help="The path to the model configuration"
|
|
)
|
|
parser.add_argument(
|
|
"--optim_config", type=str, help="The path to the optimizer configuration"
|
|
)
|
|
parser.add_argument("--split_path", type=str, help="The split file path.")
|
|
parser.add_argument("--search_shape", type=str, help="The shape to be searched.")
|
|
# parser.add_argument('--arch_para_pure', type=int, help='The architecture-parameter pure or not.')
|
|
parser.add_argument(
|
|
"--gumbel_tau_max", type=float, help="The maximum tau for Gumbel."
|
|
)
|
|
parser.add_argument(
|
|
"--gumbel_tau_min", type=float, help="The minimum tau for Gumbel."
|
|
)
|
|
parser.add_argument("--procedure", type=str, help="The procedure basic prefix.")
|
|
parser.add_argument("--FLOP_ratio", type=float, help="The expected FLOP ratio.")
|
|
parser.add_argument("--FLOP_weight", type=float, help="The loss weight for FLOP.")
|
|
parser.add_argument(
|
|
"--FLOP_tolerant", type=float, help="The tolerant range for FLOP."
|
|
)
|
|
add_shared_args(parser)
|
|
# Optimization options
|
|
parser.add_argument(
|
|
"--batch_size", type=int, default=2, help="Batch size for training."
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if args.rand_seed is None or args.rand_seed < 0:
|
|
args.rand_seed = random.randint(1, 100000)
|
|
assert args.save_dir is not None, "save-path argument can not be None"
|
|
assert args.gumbel_tau_max is not None and args.gumbel_tau_min is not None
|
|
assert (
|
|
args.FLOP_tolerant is not None and args.FLOP_tolerant > 0
|
|
), "invalid FLOP_tolerant : {:}".format(FLOP_tolerant)
|
|
# assert args.arch_para_pure is not None, 'arch_para_pure is not None: {:}'.format(args.arch_para_pure)
|
|
# args.arch_para_pure = bool(args.arch_para_pure)
|
|
return args
|