44 lines
2.8 KiB
Python
44 lines
2.8 KiB
Python
###########################################################################################
|
|
# Copyright (c) Hayeon Lee, Eunyoung Hyung [GitHub MetaD2A], 2021
|
|
# Rapid Neural Architecture Search by Learning to Generate Graphs from Datasets, ICLR 2021
|
|
###########################################################################################
|
|
import argparse
|
|
|
|
def str2bool(v):
|
|
return v.lower() in ['t', 'true', True]
|
|
|
|
|
|
def get_parser():
|
|
parser = argparse.ArgumentParser()
|
|
# general settings
|
|
parser.add_argument('--seed', type=int, default=333)
|
|
parser.add_argument('--gpu', type=str, default='0', help='set visible gpus')
|
|
parser.add_argument('--model_name', type=str, default=None, choices=['generator', 'predictor', 'train_arch'])
|
|
parser.add_argument('--save-path', type=str, default='results', help='the path of save directory')
|
|
parser.add_argument('--data-path', type=str, default='data', help='the path of save directory')
|
|
parser.add_argument('--save-epoch', type=int, default=20, help='how many epochs to wait each time to save model states')
|
|
parser.add_argument('--max-epoch', type=int, default=400, help='number of epochs to train')
|
|
parser.add_argument('--batch_size', type=int, default=32, help='batch size for generator')
|
|
parser.add_argument('--graph-data-name', default='ofa_mbv3', help='graph dataset name')
|
|
parser.add_argument('--nvt', type=int, default=27, help='number of different node types, 21 for ofa_mbv3 without in/out node')
|
|
# set encoder
|
|
parser.add_argument('--num-sample', type=int, default=20, help='the number of images as input for set encoder')
|
|
# graph encoder
|
|
parser.add_argument('--hs', type=int, default=56, help='hidden size of GRUs')
|
|
parser.add_argument('--nz', type=int, default=56, help='the number of dimensions of latent vectors z')
|
|
# test
|
|
parser.add_argument('--test', action='store_true', default=False, help='turn on test mode')
|
|
parser.add_argument('--load-epoch', type=int, default=20, help='checkpoint epoch loaded for meta-test')
|
|
parser.add_argument('--data-name', type=str, default=None, help='meta-test dataset name')
|
|
parser.add_argument('--num-class', type=int, default=None, help='the number of class of dataset')
|
|
parser.add_argument('--num-gen-arch', type=int, default=200, help='the number of candidate architectures generated by the generator')
|
|
parser.add_argument('--train-arch', type=str2bool, default=True, help='whether to train the searched architecture')
|
|
# database
|
|
parser.add_argument('--index', type=int, default=None, help='the process number when creating DB')
|
|
parser.add_argument('--imgnet', type=str, default=None, help='The path of imagenet')
|
|
parser.add_argument('--collect', action='store_true', default=False, help='whether to train the searched architecture')
|
|
|
|
args = parser.parse_args()
|
|
|
|
return args
|