Create NATS

exps/algos-v2/regularized_ea.py

@@ -23,8 +23,8 @@ from datasets     import get_datasets, SearchDataset
 from procedures   import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint, get_optim_scheduler
 from utils        import get_model_infos, obtain_accuracy
 from log_utils    import AverageMeter, time_string, convert_secs2time
-from nas_201_api  import NASBench201API, NASBench301API
 from models       import CellStructure, get_search_spaces
+from nats_bench   import create
 
 
 class Model(object):
@@ -38,47 +38,6 @@ class Model(object):
     return '{:}'.format(self.arch)
   
 
-# This function is to mimic the training and evaluatinig procedure for a single architecture `arch`.
-# The time_cost is calculated as the total training time for a few (e.g., 12 epochs) plus the evaluation time for one epoch.
-# For use_012_epoch_training = True, the architecture is trained for 12 epochs, with LR being decaded from 0.1 to 0.
-#       In this case, the LR schedular is converged.
-# For use_012_epoch_training = False, the architecture is planed to be trained for 200 epochs, but we early stop its procedure.
-#       
-def train_and_eval(arch, nas_bench, extra_info, dataname='cifar10-valid', use_012_epoch_training=True):
-
-  if use_012_epoch_training and nas_bench is not None:
-    arch_index = nas_bench.query_index_by_arch( arch )
-    assert arch_index >= 0, 'can not find this arch : {:}'.format(arch)
-    valid_acc, time_cost = info['valid-accuracy'], info['train-all-time'] + info['valid-per-time']
-    #_, valid_acc = info.get_metrics('cifar10-valid', 'x-valid' , 25, True) # use the validation accuracy after 25 training epochs
-  elif not use_012_epoch_training and nas_bench is not None:
-    # Please contact me if you want to use the following logic, because it has some potential issues.
-    # Please use `use_012_epoch_training=False` for cifar10 only.
-    # It did return values for cifar100 and ImageNet16-120, but it has some potential issues. (Please email me for more details)
-    arch_index, nepoch = nas_bench.query_index_by_arch( arch ), 25
-    assert arch_index >= 0, 'can not find this arch : {:}'.format(arch)
-    xoinfo = nas_bench.get_more_info(arch_index, 'cifar10-valid', iepoch=None, hp='12')
-    xocost = nas_bench.get_cost_info(arch_index, 'cifar10-valid', hp='200')
-    info = nas_bench.get_more_info(arch_index, dataname, nepoch, hp='200', is_random=True) # use the validation accuracy after 25 training epochs, which is used in our ICLR submission (not the camera ready).
-    cost = nas_bench.get_cost_info(arch_index, dataname, hp='200')
-    # The following codes are used to estimate the time cost.
-    # When we build NAS-Bench-201, architectures are trained on different machines and we can not use that time record.
-    # When we create checkpoints for converged_LR, we run all experiments on 1080Ti, and thus the time for each architecture can be fairly compared.
-    nums = {'ImageNet16-120-train': 151700, 'ImageNet16-120-valid': 3000,
-            'cifar10-valid-train' : 25000,  'cifar10-valid-valid' : 25000,
-            'cifar100-train'      : 50000,  'cifar100-valid'      : 5000}
-    estimated_train_cost = xoinfo['train-per-time'] / nums['cifar10-valid-train'] * nums['{:}-train'.format(dataname)] / xocost['latency'] * cost['latency'] * nepoch
-    estimated_valid_cost = xoinfo['valid-per-time'] / nums['cifar10-valid-valid'] * nums['{:}-valid'.format(dataname)] / xocost['latency'] * cost['latency']
-    try:
-      valid_acc, time_cost = info['valid-accuracy'], estimated_train_cost + estimated_valid_cost
-    except:
-      valid_acc, time_cost = info['valtest-accuracy'], estimated_train_cost + estimated_valid_cost
-  else:
-    # train a model from scratch.
-    raise ValueError('NOT IMPLEMENT YET')
-  return valid_acc, time_cost
-
-
 def random_topology_func(op_names, max_nodes=4):
   # Return a random architecture
   def random_architecture():
@@ -239,12 +198,7 @@ if __name__ == '__main__':
   parser.add_argument('--rand_seed',          type=int,   default=-1,    help='manual seed')
   args = parser.parse_args()
 
-  if args.search_space == 'tss':
-    api = NASBench201API(verbose=False)
-  elif args.search_space == 'sss':
-    api = NASBench301API(verbose=False)
-  else:
-    raise ValueError('Invalid search space : {:}'.format(args.search_space))
+  api = create(None, args.search_space, verbose=False)
 
   args.save_dir = os.path.join('{:}-{:}'.format(args.save_dir, args.search_space), args.dataset, 'R-EA-SS{:}'.format(args.ea_sample_size))
   print('save-dir : {:}'.format(args.save_dir))