import os import sys import random import numpy as np import argparse import torch import os from nag import NAG # sys.path.append(os.getcwd()) # from utils import str2bool def str2bool(v): return v.lower() in ['t', 'true', True] # save_path = "results" # data_path = os.path.join('MetaD2A_nas_bench_201', 'data') # model_load_path = '/home/data/GTAD/baselines/transferNAS' def get_parser(): parser = argparse.ArgumentParser() # general settings parser.add_argument('--seed', type=int, default=444) parser.add_argument('--gpu', type=str, default='0', help='set visible gpus') parser.add_argument('--search_space', type=str, default='ofa') parser.add_argument('--save-path', type=str, default=None, help='the path of save directory') parser.add_argument('--data-path', type=str, default=None, help='the path of save directory') parser.add_argument('--model-load-path', type=str, default=None, help='') 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=50, help='number of epochs to train') parser.add_argument('--batch_size', type=int, default=1024, help='batch size for generator') parser.add_argument('--graph-data-name', default='ofa', help='graph dataset name') parser.add_argument('--nvt', type=int, default=27, help='number of different node types') # 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=512, 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=True, help='turn on test mode') parser.add_argument('--load-epoch', type=int, default=100, help='checkpoint epoch loaded for meta-test') parser.add_argument('--data-name', type=str, default='pets', help='meta-test dataset name') parser.add_argument('--trials', type=int, default=5) 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=500, 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') parser.add_argument('--n_training_samples', type=int, default=5) parser.add_argument('--N', type=int, default=10) parser.add_argument('--use_gp', type=str2bool, default=False) parser.add_argument('--sorting', type=str2bool, default=True) parser.add_argument('--use_metad2a_predictor_selec', type=str2bool, default=True) parser.add_argument('--use_ensemble_selec', type=str2bool, default=False) # ---------- For diffusion NAG ------------ # parser.add_argument('--folder_name', type=str, default='DiffusionNAG') parser.add_argument('--task', type=str, default='mtst') parser.add_argument('--exp_name', type=str, default='') parser.add_argument('--wandb_exp_name', type=str, default='') parser.add_argument('--wandb_project_name', type=str, default='DiffusionNAG') parser.add_argument('--use_wandb', type=str2bool, default=False) parser.add_argument('--classifier_scale', type=int, default=10000.0, help='classifier scale') parser.add_argument('--eval_batch_size', type=int, default=256) parser.add_argument('--predictor', type=str, default='euler_maruyama', choices=['euler_maruyama', 'reverse_diffusion', 'none']) parser.add_argument('--corrector', type=str, default='langevin', choices=['none', 'langevin']) parser.add_argument('--weight_ratio', type=str2bool, default=False) parser.add_argument('--weight_scheduling', type=str2bool, default=False) parser.add_argument('--weight_ratio_abs', type=str2bool, default=False) parser.add_argument('--p', type=int, default=1) parser.add_argument('--prod_w', type=str2bool, default=False) parser.add_argument('--t_spot', type=float, default=1.0) parser.add_argument('--t_spot_end', type=float, default=0.0) # Train parser.add_argument('--lr', type=float, default=0.001, help='learning rate') parser.add_argument('--epochs', type=int, default=500) parser.add_argument('--save_best_p_corr', type=str2bool, default=True) parser.add_argument('--unique', type=str2bool, default=True) parser.add_argument('--patient_factor', type=int, default=20) parser.add_argument('--n_gen_samples', type=int, default=50) ################ OFA #################### parser.add_argument('--ofa_path', type=str, default='/home/hayeon/imagenet1k', help='') parser.add_argument('--ofa_batch_size', type=int, default=256, help='') parser.add_argument('--ofa_workers', type=int, default=4, help='') ################ Diffusion ############## parser.add_argument('--diffusion_lr', type=float, default=1e-3, help='') parser.add_argument('--noise_aware_acc_norm', type=int, default=-1) parser.add_argument('--fix_task', type=str2bool, default=True) ################ BO #################### parser.add_argument('--bo_loop_max_epoch', type=int, default=30) parser.add_argument('--bo_loop_acc_norm', type=int, default=1) parser.add_argument('--gp_model_acc_norm', type=int, default=1) parser.add_argument('--num_ensemble', type=int, default=3) parser.add_argument('--explore_type', type=str, default='ei') ################ BO #################### # parser.add_argument('--multi_proc', type=str2bool, default=False) parser.add_argument('--eps', type=float, default=0.) parser.add_argument('--beta', type=float, default=0.5) parser.add_argument('--pos_enc_type', type=int, default=4) args = parser.parse_args() return args def set_exp_name(args): exp_name = f'./exp/{args.task}/{args.folder_name}/data-{args.data_name}' wandb_exp_name = f'./exp/{args.task}/{args.folder_name}/{args.data_name}' os.makedirs(exp_name, exist_ok=True) args.exp_name = exp_name args.wandb_exp_name = wandb_exp_name def main(): args = get_parser() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") torch.cuda.manual_seed(args.seed) torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) set_exp_name(args) p = NAG(args) if args.test: p.meta_test() else: p.meta_train() if __name__ == '__main__': main()