155 lines
7.0 KiB
Python
155 lines
7.0 KiB
Python
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()
|