first commit

NAS-Bench-201/main_exp/transfer_nag/main.py

@@ -0,0 +1,91 @@
+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())
+save_path = "results"
+data_path = os.path.join('MetaD2A_nas_bench_201', 'data')
+
+
+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=444)
+    parser.add_argument('--gpu', type=str, default='0', help='set visible gpus')
+    parser.add_argument('--save-path', type=str, default=save_path, help='the path of save directory')
+    parser.add_argument('--data-path', type=str, default=data_path, help='the path of save directory')
+    parser.add_argument('--model-load-path', type=str, default='', 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=1000, 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='nasbench201', help='graph dataset name')
+    parser.add_argument('--nvt', type=int, default=7, help='number of different node types, 7: NAS-Bench-201 including 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=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=20)
+    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_init', type=int, default=10)
+    parser.add_argument('--N', type=int, default=1)
+    # DiffusionNAG
+    parser.add_argument('--folder_name', type=str, default='debug')
+    parser.add_argument('--exp_name', type=str, default='')
+    parser.add_argument('--classifier_scale', type=float, default=10000., help='classifier scale')
+    parser.add_argument('--scale_step', type=float, default=300.)
+    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('--patient_factor', type=int, default=20)
+    parser.add_argument('--n_gen_samples', type=int, default=10)
+    parser.add_argument('--multi_proc', type=str2bool, default=True)
+
+    args = parser.parse_args()
+    return args
+
+
+def set_exp_name(args):
+    exp_name = f'./results/transfer_nag/{args.folder_name}/{args.data_name}'
+    os.makedirs(exp_name, exist_ok=True)
+    args.exp_name = exp_name
+
+
+def main():
+    ## Get arguments
+    args = get_parser()
+
+    ## Set gpus and seeds
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+    torch.cuda.manual_seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    np.random.seed(args.seed)
+    random.seed(args.seed)
+
+    ## Set experiment name
+    set_exp_name(args)
+
+    ## Run
+    nag = NAG(args)
+    nag.meta_test()
+
+
+if __name__ == '__main__':
+    main()