Create NATS

CHANGE-LOG.md

@@ -5,3 +5,4 @@
 - [2019.09.28] [f8f3f38] TAS and SETN codes were publicly released.
 - [2019.01.31] [13e908f] GDAS codes were publicly released.
 - [2020.07.01] [a45808b] Upgrade NAS-API to the 2.0 version.
+- [2020.07.30] [       ] Create NATS-BENCH.

exps/NAS-Bench-201/test-nas-api.py

@@ -1,9 +1,11 @@
 ###############################################################
 # NAS-Bench-201, ICLR 2020 (https://arxiv.org/abs/2001.00326) #
 ###############################################################
+# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size
+###############################################################
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.06           #
 ###############################################################
-# Usage: python exps/NAS-Bench-201/test-nas-api.py
+# Usage: python exps/NAS-Bench-201/test-nas-api.py            #
 ###############################################################
 import os, sys, time, torch, argparse
 import numpy as np
@@ -21,7 +23,7 @@ import matplotlib.ticker as ticker
 lib_dir = (Path(__file__).parent / '..' / '..' / 'lib').resolve()
 if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
 from config_utils import dict2config, load_config
-from nas_201_api import NASBench201API, NASBench301API
+from nats_bench import create
 from log_utils import time_string
 from models import get_cell_based_tiny_net, CellStructure
 
@@ -97,15 +99,14 @@ def test_issue_81_82(api):
 
 if __name__ == '__main__':
 
-  api201 = NASBench201API(os.path.join(os.environ['TORCH_HOME'], 'NAS-Bench-201-v1_0-e61699.pth'), verbose=True)
+  api201 = create(os.path.join(os.environ['TORCH_HOME'], 'NAS-Bench-201-v1_0-e61699.pth'), 'topology', True)
   test_issue_81_82(api201)
-  # test_api(api201, False)
   print ('Test {:} done'.format(api201))
 
-  api201 = NASBench201API(None, verbose=True)
+  api201 = create(None, 'topology', True)  # use the default file path
   test_issue_81_82(api201)
   test_api(api201, False)
   print ('Test {:} done'.format(api201))
 
-  # api301 = NASBench301API(None, verbose=True)
+  api301 = create(None, 'size', True)
-  # test_api(api301, True)
+  test_api(api301, True)

exps/NAS-Bench-201/xshape-collect.py

@@ -16,7 +16,7 @@ from log_utils    import AverageMeter, time_string, convert_secs2time
 from config_utils import dict2config
 # NAS-Bench-201 related module or function
 from models       import CellStructure, get_cell_based_tiny_net
-from nas_201_api  import NASBench301API, ArchResults, ResultsCount
+from nas_201_api  import ArchResults, ResultsCount
 from procedures   import bench_pure_evaluate as pure_evaluate, get_nas_bench_loaders
 
 

exps/algos-v2/README.md

@@ -1 +1,3 @@
-# Benchmarking NAS Algorithms
+# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size
+
+# Benchmarking 13 NAS Algorithm

exps/algos-v2/bohb.py

@@ -18,7 +18,7 @@ from config_utils import load_config
 from datasets     import get_datasets, SearchDataset
 from procedures   import prepare_seed, prepare_logger
 from log_utils    import AverageMeter, time_string, convert_secs2time
-from nas_201_api  import NASBench201API, NASBench301API
+from nats_bench   import create
 from models       import CellStructure, get_search_spaces
 # BOHB: Robust and Efficient Hyperparameter Optimization at Scale, ICML 2018
 import ConfigSpace
@@ -167,12 +167,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 = create(None, args.search_space, verbose=False)
-    api = NASBench201API(verbose=False)
-  elif args.search_space == 'sss':
-    api = NASBench301API(verbose=False)
-  else:
-    raise ValueError('Invalid search space : {:}'.format(args.search_space))
 
   args.save_dir = os.path.join('{:}-{:}'.format(args.save_dir, args.search_space), args.dataset, 'BOHB')
   print('save-dir : {:}'.format(args.save_dir))

exps/algos-v2/random_wo_share.py

@@ -21,7 +21,7 @@ from procedures   import prepare_seed, prepare_logger, save_checkpoint, copy_che
 from utils        import get_model_infos, obtain_accuracy
 from log_utils    import AverageMeter, time_string, convert_secs2time
 from models       import get_search_spaces
-from nas_201_api  import NASBench201API, NASBench301API
+from nats_bench   import create
 from regularized_ea import random_topology_func, random_size_func
 
 
@@ -71,12 +71,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 = create(None, args.search_space, verbose=False)
-    api = NASBench201API(verbose=False)
-  elif args.search_space == 'sss':
-    api = NASBench301API(verbose=False)
-  else:
-    raise ValueError('Invalid search space : {:}'.format(args.search_space))
 
   args.save_dir = os.path.join('{:}-{:}'.format(args.save_dir, args.search_space), args.dataset, 'RANDOM')
   print('save-dir : {:}'.format(args.save_dir))

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 = create(None, args.search_space, verbose=False)
-    api = NASBench201API(verbose=False)
-  elif args.search_space == 'sss':
-    api = NASBench301API(verbose=False)
-  else:
-    raise ValueError('Invalid search space : {:}'.format(args.search_space))
 
   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))

exps/algos-v2/reinforce.py

@@ -24,8 +24,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 PolicyTopology(nn.Module):
@@ -192,12 +192,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 = create(None, args.search_space, verbose=False)
-    api = NASBench201API(verbose=False)
-  elif args.search_space == 'sss':
-    api = NASBench301API(verbose=False)
-  else:
-    raise ValueError('Invalid search space : {:}'.format(args.search_space))
 
   args.save_dir = os.path.join('{:}-{:}'.format(args.save_dir, args.search_space), args.dataset, 'REINFORCE-{:}'.format(args.learning_rate))
   print('save-dir : {:}'.format(args.save_dir))

exps/algos-v2/search-cell.py

@@ -39,7 +39,7 @@ from procedures   import prepare_seed, prepare_logger, save_checkpoint, copy_che
 from utils        import count_parameters_in_MB, obtain_accuracy
 from log_utils    import AverageMeter, time_string, convert_secs2time
 from models       import get_cell_based_tiny_net, get_search_spaces
-from nas_201_api  import NASBench201API as API
+from nats_bench   import create
 
 
 # The following three functions are used for DARTS-V2
@@ -364,7 +364,7 @@ def main(xargs):
   logger.log('The parameters of the search model = {:.2f} MB'.format(params))
   logger.log('search-space : {:}'.format(search_space))
   if bool(xargs.use_api):
-    api = API(verbose=False)
+    api = create(None, 'topology', verbose=False)
   else:
     api = None
   logger.log('{:} create API = {:} done'.format(time_string(), api))

exps/algos-v2/search-size.py

@@ -27,7 +27,7 @@ from procedures   import prepare_seed, prepare_logger, save_checkpoint, copy_che
 from utils        import count_parameters_in_MB, obtain_accuracy
 from log_utils    import AverageMeter, time_string, convert_secs2time
 from models       import get_cell_based_tiny_net, get_search_spaces
-from nas_201_api  import NASBench301API as API
+from nats_bench   import create
 
 
 # Ad-hoc for TuNAS
@@ -176,7 +176,7 @@ def main(xargs):
   logger.log('The parameters of the search model = {:.2f} MB'.format(params))
   logger.log('search-space : {:}'.format(search_space))
   if bool(xargs.use_api):
-    api = API(verbose=False)
+    api = create(None, 'size', verbose=False)
   else:
     api = None
   logger.log('{:} create API = {:} done'.format(time_string(), api))
@@ -291,7 +291,7 @@ if __name__ == '__main__':
   parser.add_argument('--rand_seed',          type=int,   help='manual seed')
   args = parser.parse_args()
   if args.rand_seed is None or args.rand_seed < 0: args.rand_seed = random.randint(1, 100000)
-  dirname = '{:}-affine{:}_BN{:}'.format(args.algo, args.affine, args.track_running_stats)
+  dirname = '{:}-affine{:}_BN{:}-AWD{:}'.format(args.algo, args.affine, args.track_running_stats, args.arch_weight_decay)
   if args.overwite_epochs is not None:
     dirname = dirname + '-E{:}'.format(args.overwite_epochs)
   args.save_dir = os.path.join('{:}-{:}'.format(args.save_dir, args.search_space), args.dataset, dirname)

exps/experimental/example-nas-bench.py

@@ -16,7 +16,7 @@ matplotlib.use('agg')
 import matplotlib.pyplot as plt
 lib_dir = (Path(__file__).parent / '..' / '..' / 'lib').resolve()
 if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
-from nas_201_api import NASBench201API, NASBench301API
+from nas_201_api import NASBench201API
 from log_utils import time_string
 from models import get_cell_based_tiny_net
 from utils import weight_watcher

exps/experimental/test-ww-bench.py

@@ -3,9 +3,6 @@
 ###########################################################################################################################################################
 # Before run these commands, the files must be properly put.
 #
-# python exps/experimental/test-ww-bench.py --base_path $HOME/.torch/NAS-Bench-201-v1_0-e61699
-# python exps/experimental/test-ww-bench.py --base_path $HOME/.torch/NAS-Bench-201-v1_1-096897 --dataset cifar10-valid --use_12 1 --use_valid 1
-# CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=4 python exps/experimental/test-ww-bench.py --base_path $HOME/.torch/NAS-Bench-201-v1_1-096897 --dataset cifar10
 # CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=4 python exps/experimental/test-ww-bench.py --search_space sss --base_path $HOME/.torch/NAS-Bench-301-v1_0 --dataset cifar10
 # CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=4 python exps/experimental/test-ww-bench.py --search_space sss --base_path $HOME/.torch/NAS-Bench-301-v1_0 --dataset cifar100
 # CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=4 python exps/experimental/test-ww-bench.py --search_space sss --base_path $HOME/.torch/NAS-Bench-301-v1_0 --dataset ImageNet16-120
@@ -22,8 +19,8 @@ matplotlib.use('agg')
 import matplotlib.pyplot as plt
 lib_dir = (Path(__file__).parent / '..' / '..' / 'lib').resolve()
 if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
-from nas_201_api import NASBench201API, NASBench301API
 from log_utils import time_string
+from nats_bench import create
 from models import get_cell_based_tiny_net
 from utils import weight_watcher
 
@@ -52,8 +49,8 @@ def evaluate(api, weight_dir, data: str):
     # compute the weight watcher results
     config = api.get_net_config(arch_index, data)
     net = get_cell_based_tiny_net(config)
-    meta_info = api.query_meta_info_by_index(arch_index, hp='200' if isinstance(api, NASBench201API) else '90')
+    meta_info = api.query_meta_info_by_index(arch_index, hp='200' if api.search_space_name == 'topology' else '90')
-    params = meta_info.get_net_param(data, 888 if isinstance(api, NASBench201API) else 777)
+    params = meta_info.get_net_param(data, 888 if api.search_space_name == 'topology' else 777)
     with torch.no_grad():
       net.load_state_dict(params)
       _, summary = weight_watcher.analyze(net, alphas=False)
@@ -70,7 +67,7 @@ def evaluate(api, weight_dir, data: str):
       ok += 1
       norms.append(cur_norm)
     # query the accuracy
-    info = meta_info.get_metrics(data, 'ori-test', iepoch=None, is_random=888 if isinstance(api, NASBench201API) else 777)
+    info = meta_info.get_metrics(data, 'ori-test', iepoch=None, is_random=888 if api.search_space_name == 'topology' else 777)
     accuracies.append(info['accuracy'])
     del net, meta_info
     # print the information
@@ -81,9 +78,8 @@ def evaluate(api, weight_dir, data: str):
 
 
 def main(search_space, meta_file: str, weight_dir, save_dir, xdata):
-  API = NASBench201API if search_space == 'tss' else NASBench301API
   save_dir.mkdir(parents=True, exist_ok=True)
-  api = API(meta_file, verbose=False)
+  api = create(meta_file, search_space, verbose=False)
   datasets = ['cifar10-valid', 'cifar10', 'cifar100', 'ImageNet16-120']
   print(time_string() + ' ' + '='*50)
   for data in datasets:

exps/experimental/vis-nats-bench-algos.py

@@ -3,8 +3,8 @@
 ###############################################################
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.06           #
 ###############################################################
-# Usage: python exps/experimental/vis-bench-algos.py --search_space tss
+# Usage: python exps/experimental/vis-nats-bench-algos.py --search_space tss
-# Usage: python exps/experimental/vis-bench-algos.py --search_space sss
+# Usage: python exps/experimental/vis-nats-bench-algos.py --search_space sss
 ###############################################################
 import os, gc, sys, time, torch, argparse
 import numpy as np
@@ -22,7 +22,7 @@ import matplotlib.ticker as ticker
 lib_dir = (Path(__file__).parent / '..' / '..' / 'lib').resolve()
 if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
 from config_utils import dict2config, load_config
-from nas_201_api import NASBench201API, NASBench301API
+from nats_bench import create
 from log_utils import time_string
 
 
@@ -48,18 +48,19 @@ def fetch_data(root_dir='./output/search', search_space='tss', dataset=None):
 
 
 def query_performance(api, data, dataset, ticket):
-  results, is_301 = [], isinstance(api, NASBench301API)
+  results, is_size_space = [], api.search_space_name == 'size'
   for i, info in data.items():
     time_w_arch = sorted(info['time_w_arch'], key=lambda x: abs(x[0]-ticket))
     time_a, arch_a = time_w_arch[0]
     time_b, arch_b = time_w_arch[1]
-    info_a = api.get_more_info(arch_a, dataset=dataset, hp=90 if is_301 else 200, is_random=False)
+    info_a = api.get_more_info(arch_a, dataset=dataset, hp=90 if is_size_space else 200, is_random=False)
-    info_b = api.get_more_info(arch_b, dataset=dataset, hp=90 if is_301 else 200, is_random=False)
+    info_b = api.get_more_info(arch_b, dataset=dataset, hp=90 if is_size_space else 200, is_random=False)
     accuracy_a, accuracy_b = info_a['test-accuracy'], info_b['test-accuracy']
     interplate = (time_b-ticket) / (time_b-time_a) * accuracy_a + (ticket-time_a) / (time_b-time_a) * accuracy_b
     results.append(interplate)
   return sum(results) / len(results)
 
+
 y_min_s = {('cifar10', 'tss'): 90,
            ('cifar10', 'sss'): 92,
            ('cifar100', 'tss'): 65,
@@ -74,6 +75,10 @@ y_max_s = {('cifar10', 'tss'): 94.5,
            ('ImageNet16-120', 'tss'): 44,
            ('ImageNet16-120', 'sss'): 46}
 
+name2label = {'cifar10': 'CIFAR-10',
+              'cifar100': 'CIFAR-100',
+              'ImageNet16-120': 'ImageNet-16-120'}
+
 def visualize_curve(api, vis_save_dir, search_space, max_time):
   vis_save_dir = vis_save_dir.resolve()
   vis_save_dir.mkdir(parents=True, exist_ok=True)
@@ -99,8 +104,8 @@ def visualize_curve(api, vis_save_dir, search_space, max_time):
       alg2accuracies[alg] = accuracies
       ax.plot([x/100 for x in time_tickets], accuracies, c=colors[idx], label='{:}'.format(alg))
       ax.set_xlabel('Estimated wall-clock time (1e2 seconds)', fontsize=LabelSize)
-      ax.set_ylabel('Test accuracy on {:}'.format(dataset), fontsize=LabelSize)
+      ax.set_ylabel('Test accuracy on {:}'.format(name2label[dataset]), fontsize=LabelSize)
-      ax.set_title('Searching results on {:}'.format(dataset), fontsize=LabelSize+4)
+      ax.set_title('Searching results on {:}'.format(name2label[dataset]), fontsize=LabelSize+4)
     ax.legend(loc=4, fontsize=LegendFontsize)
 
   fig, axs = plt.subplots(1, 3, figsize=figsize)
@@ -123,10 +128,5 @@ if __name__ == '__main__':
 
   save_dir = Path(args.save_dir)
 
-  if args.search_space == 'tss':
+  api = create(None, args.search_space, verbose=False)
-    api = NASBench201API(verbose=False)
-  elif args.search_space == 'sss':
-    api = NASBench301API(verbose=False)
-  else:
-    raise ValueError('Invalid search space : {:}'.format(args.search_space))
   visualize_curve(api, save_dir, args.search_space, args.max_time)

exps/experimental/vis-nats-bench-ws.py

@@ -3,8 +3,8 @@
 ###############################################################
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.06           #
 ###############################################################
-# Usage: python exps/experimental/vis-bench-ws.py --search_space tss
+# Usage: python exps/experimental/vis-nats-bench-ws.py --search_space tss
-# Usage: python exps/experimental/vis-bench-ws.py --search_space sss
+# Usage: python exps/experimental/vis-nats-bench-ws.py --search_space sss
 ###############################################################
 import os, gc, sys, time, torch, argparse
 import numpy as np
@@ -22,15 +22,16 @@ import matplotlib.ticker as ticker
 lib_dir = (Path(__file__).parent / '..' / '..' / 'lib').resolve()
 if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
 from config_utils import dict2config, load_config
-from nas_201_api import NASBench201API, NASBench301API
+from nats_bench import create
 from log_utils import time_string
 
 
 def fetch_data(root_dir='./output/search', search_space='tss', dataset=None):
   ss_dir = '{:}-{:}'.format(root_dir, search_space)
   alg2name, alg2path = OrderedDict(), OrderedDict()
+  seeds = [777, 888, 999]
+  print('\n[fetch data] from {:} on {:}'.format(search_space, dataset))
   if search_space == 'tss':
-    seeds = [777]
     alg2name['GDAS'] = 'gdas-affine0_BN0-None'
     alg2name['RSPS'] = 'random-affine0_BN0-None'
     alg2name['DARTS (1st)'] = 'darts-v1-affine0_BN0-None'
@@ -38,7 +39,6 @@ def fetch_data(root_dir='./output/search', search_space='tss', dataset=None):
     alg2name['ENAS'] = 'enas-affine0_BN0-None'
     alg2name['SETN'] = 'setn-affine0_BN0-None'
   else:
-    seeds = [777, 888, 999]
     alg2name['TAS'] = 'tas-affine0_BN0'
     alg2name['FBNetV2'] = 'fbv2-affine0_BN0'
     alg2name['TuNAS'] = 'tunas-affine0_BN0'
@@ -46,13 +46,19 @@ def fetch_data(root_dir='./output/search', search_space='tss', dataset=None):
     alg2path[alg] = os.path.join(ss_dir, dataset, name, 'seed-{:}-last-info.pth')
   alg2data = OrderedDict()
   for alg, path in alg2path.items():
-    alg2data[alg] = []
+    alg2data[alg], ok_num = [], 0
     for seed in seeds:
       xpath = path.format(seed)
-      assert os.path.isfile(xpath), 'invalid path : {:}'.format(xpath)
+      if os.path.isfile(xpath):
+        ok_num += 1
+      else:
+        print('This is an invalid path : {:}'.format(xpath))
+        continue
       data = torch.load(xpath, map_location=torch.device('cpu'))
       data = torch.load(data['last_checkpoint'], map_location=torch.device('cpu'))
       alg2data[alg].append(data['genotypes'])
+    print('This algorithm : {:} has {:} valid ckps.'.format(alg, ok_num))
+    assert ok_num > 0, 'Must have at least 1 valid ckps.'
   return alg2data
 
 
@@ -95,7 +101,7 @@ def visualize_curve(api, vis_save_dir, search_space):
       for iepoch in range(epochs+1):
         structures, accs = [_[iepoch-1] for _ in data], []
         for structure in structures:
-          info = api.get_more_info(structure, dataset=dataset, hp=90 if isinstance(api, NASBench301API) else 200, is_random=False)
+          info = api.get_more_info(structure, dataset=dataset, hp=90 if api.search_space_name == 'size' else 200, is_random=False)
           accs.append(info['test-accuracy'])
         accuracies.append(sum(accs)/len(accs))
         xs.append(iepoch)
@@ -124,12 +130,6 @@ if __name__ == '__main__':
   args = parser.parse_args()
 
   save_dir = Path(args.save_dir)
-  alg2data = fetch_data(search_space='tss', dataset='cifar10')
 
-  if args.search_space == 'tss':
+  api = create(None, args.search_space, verbose=False)
-    api = NASBench201API(verbose=False)
-  elif args.search_space == 'sss':
-    api = NASBench301API(verbose=False)
-  else:
-    raise ValueError('Invalid search space : {:}'.format(args.search_space))
   visualize_curve(api, save_dir, args.search_space)

exps/experimental/visualize-nas-bench-x.py

@@ -21,9 +21,9 @@ import matplotlib.ticker as ticker
 lib_dir = (Path(__file__).parent / '..' / '..' / 'lib').resolve()
 if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
 from config_utils import dict2config, load_config
-from nas_201_api import NASBench201API, NASBench301API
 from log_utils import time_string
 from models import get_cell_based_tiny_net
+from nats_bench import create
 
 
 def visualize_info(api, vis_save_dir, indicator):
@@ -391,11 +391,11 @@ if __name__ == '__main__':
   to_save_dir = Path(args.save_dir)
 
   datasets = ['cifar10', 'cifar100', 'ImageNet16-120']
-  api201 = NASBench201API(None, verbose=True)
+  api201 = create(None, 'tss', verbose=True)
   for xdata in datasets:
     visualize_tss_info(api201, xdata, to_save_dir)
 
-  api301 = NASBench301API(None, verbose=True)
+  api301 = create(None, 'size', verbose=True)
   for xdata in datasets:
     visualize_sss_info(api301, xdata, to_save_dir)
 

lib/models/__init__.py

@@ -64,7 +64,7 @@ def get_search_spaces(xtype, name) -> List[Text]:
     assert name in SearchSpaceNames, 'invalid name [{:}] in {:}'.format(name, SearchSpaceNames.keys())
     return SearchSpaceNames[name]
   elif xtype == 'sss':  # The size search space.
-    if name == 'nas-bench-301':
+    if name == 'nas-bench-301' or name == 'nats-bench' or name == 'nats-bench-size':
       return {'candidates': [8, 16, 24, 32, 40, 48, 56, 64],
               'numbers': 5}
     else:

lib/models/cell_operations.py

@@ -25,6 +25,7 @@ NAS_BENCH_201         = ['none', 'skip_connect', 'nor_conv_1x1', 'nor_conv_3x3',
 DARTS_SPACE           = ['none', 'skip_connect', 'dua_sepc_3x3', 'dua_sepc_5x5', 'dil_sepc_3x3', 'dil_sepc_5x5', 'avg_pool_3x3', 'max_pool_3x3']
 
 SearchSpaceNames = {'connect-nas'  : CONNECT_NAS_BENCHMARK,
+                    'nats-bench'   : NAS_BENCH_201,
                     'nas-bench-201': NAS_BENCH_201,
                     'nas-bench-301': NAS_BENCH_201,
                     'darts'        : DARTS_SPACE}

lib/nats_bench/__init__.py

@@ -0,0 +1,25 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.07 #
+#####################################################
+# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size
+#####################################################
+#
+#
+from .api_utils import ArchResults, ResultsCount
+from .api_topology import NATStopology
+from .api_size import NATSsize
+
+NATS_BENCH_API_VERSIONs = ['v1.0']    # [2020.07.30]
+
+
+def version():
+  return NATS_BENCH_API_VERSIONs[-1]
+
+
+def create(file_path_or_dict, search_space, verbose=True):
+  if search_space in ['tss', 'topology']:
+    return NATStopology(file_path_or_dict, verbose)
+  elif search_space in ['sss', 'size']:
+    return NATSsize(file_path_or_dict, verbose)
+  else:
+    raise ValueError('invalid search space : {:}'.format(search_space))

lib/nats_bench/api_size.py

@@ -0,0 +1,222 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.06 #
+############################################################################################
+# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size
+############################################################################################
+# The history of benchmark files:
+# 
+import os, copy, random, torch, numpy as np
+from pathlib import Path
+from typing import List, Text, Union, Dict, Optional
+from collections import OrderedDict, defaultdict
+from .api_utils import ArchResults
+from .api_utils import NASBenchMetaAPI
+from .api_utils import remap_dataset_set_names
+
+
+ALL_BENCHMARK_FILES = ['NAS-Bench-301-v1_0-363be7.pth']
+ALL_ARCHIVE_DIRS = ['NAS-Bench-301-v1_0-archive']
+
+
+def print_information(information, extra_info=None, show=False):
+  dataset_names = information.get_dataset_names()
+  strings = [information.arch_str, 'datasets : {:}, extra-info : {:}'.format(dataset_names, extra_info)]
+  def metric2str(loss, acc):
+    return 'loss = {:.3f} & top1 = {:.2f}%'.format(loss, acc)
+
+  for ida, dataset in enumerate(dataset_names):
+    metric = information.get_compute_costs(dataset)
+    flop, param, latency = metric['flops'], metric['params'], metric['latency']
+    str1 = '{:14s} FLOP={:6.2f} M, Params={:.3f} MB, latency={:} ms.'.format(dataset, flop, param, '{:.2f}'.format(latency*1000) if latency is not None and latency > 0 else None)
+    train_info = information.get_metrics(dataset, 'train')
+    if dataset == 'cifar10-valid':
+      valid_info = information.get_metrics(dataset, 'x-valid')
+      test__info = information.get_metrics(dataset, 'ori-test')
+      str2 = '{:14s} train : [{:}], valid : [{:}], test : [{:}]'.format(
+                dataset, metric2str(train_info['loss'], train_info['accuracy']),
+                metric2str(valid_info['loss'], valid_info['accuracy']),
+                metric2str(test__info['loss'], test__info['accuracy']))
+    elif dataset == 'cifar10':
+      test__info = information.get_metrics(dataset, 'ori-test')
+      str2 = '{:14s} train : [{:}], test  : [{:}]'.format(dataset, metric2str(train_info['loss'], train_info['accuracy']), metric2str(test__info['loss'], test__info['accuracy']))
+    else:
+      valid_info = information.get_metrics(dataset, 'x-valid')
+      test__info = information.get_metrics(dataset, 'x-test')
+      str2 = '{:14s} train : [{:}], valid : [{:}], test : [{:}]'.format(dataset, metric2str(train_info['loss'], train_info['accuracy']), metric2str(valid_info['loss'], valid_info['accuracy']), metric2str(test__info['loss'], test__info['accuracy']))
+    strings += [str1, str2]
+  if show: print('\n'.join(strings))
+  return strings
+
+
+"""
+This is the class for the API of size search space in NATS-Bench.
+"""
+class NATSsize(NASBenchMetaAPI):
+
+  """ The initialization function that takes the dataset file path (or a dict loaded from that path) as input. """
+  def __init__(self, file_path_or_dict: Optional[Union[Text, Dict]]=None, verbose: bool=True):
+    self.filename = None
+    self._search_space_name = 'size'
+    self.reset_time()
+    if file_path_or_dict is None:
+      file_path_or_dict = os.path.join(os.environ['TORCH_HOME'], ALL_BENCHMARK_FILES[-1])
+      print ('Try to use the default NATS-Bench (size) path from {:}.'.format(file_path_or_dict))
+    if isinstance(file_path_or_dict, str) or isinstance(file_path_or_dict, Path):
+      file_path_or_dict = str(file_path_or_dict)
+      if verbose: print('try to create the NATS-Bench (size) api from {:}'.format(file_path_or_dict))
+      assert os.path.isfile(file_path_or_dict), 'invalid path : {:}'.format(file_path_or_dict)
+      self.filename = Path(file_path_or_dict).name
+      file_path_or_dict = torch.load(file_path_or_dict, map_location='cpu')
+    elif isinstance(file_path_or_dict, dict):
+      file_path_or_dict = copy.deepcopy( file_path_or_dict )
+    else: raise ValueError('invalid type : {:} not in [str, dict]'.format(type(file_path_or_dict)))
+    assert isinstance(file_path_or_dict, dict), 'It should be a dict instead of {:}'.format(type(file_path_or_dict))
+    self.verbose = verbose # [TODO] a flag indicating whether to print more logs
+    keys = ('meta_archs', 'arch2infos', 'evaluated_indexes')
+    for key in keys: assert key in file_path_or_dict, 'Can not find key[{:}] in the dict'.format(key)
+    self.meta_archs = copy.deepcopy( file_path_or_dict['meta_archs'] )
+    # This is a dict mapping each architecture to a dict, where the key is #epochs and the value is ArchResults
+    self.arch2infos_dict = OrderedDict()
+    self._avaliable_hps = set()
+    for xkey in sorted(list(file_path_or_dict['arch2infos'].keys())):
+      all_infos = file_path_or_dict['arch2infos'][xkey]
+      hp2archres = OrderedDict()
+      for hp_key, results in all_infos.items():
+        hp2archres[hp_key] = ArchResults.create_from_state_dict(results)
+        self._avaliable_hps.add(hp_key)  # save the avaliable hyper-parameter
+      self.arch2infos_dict[xkey] = hp2archres
+    self.evaluated_indexes = sorted(list(file_path_or_dict['evaluated_indexes']))
+    self.archstr2index = {}
+    for idx, arch in enumerate(self.meta_archs):
+      assert arch not in self.archstr2index, 'This [{:}]-th arch {:} already in the dict ({:}).'.format(idx, arch, self.archstr2index[arch])
+      self.archstr2index[ arch ] = idx
+    if self.verbose:
+      print('Create NATS-Bench (size) done with {:}/{:} architectures avaliable.'.format(len(self.evaluated_indexes), len(self.meta_archs)))
+
+  def reload(self, archive_root: Text = None, index: int = None):
+    """Overwrite all information of the 'index'-th architecture in the search space, where the data will be loaded from 'archive_root'.
+       If index is None, overwrite all ckps.
+    """
+    if self.verbose:
+      print('Call clear_params with archive_root={:} and index={:}'.format(archive_root, index))
+    if archive_root is None:
+      archive_root = os.path.join(os.environ['TORCH_HOME'], ALL_ARCHIVE_DIRS[-1])
+    assert os.path.isdir(archive_root), 'invalid directory : {:}'.format(archive_root)
+    if index is None:
+      indexes = list(range(len(self)))
+    else:
+      indexes = [index]
+    for idx in indexes:
+      assert 0 <= idx < len(self.meta_archs), 'invalid index of {:}'.format(idx)
+      xfile_path = os.path.join(archive_root, '{:06d}-FULL.pth'.format(idx))
+      if not os.path.isfile(xfile_path):
+        xfile_path = os.path.join(archive_root, '{:d}-FULL.pth'.format(idx))
+      assert os.path.isfile(xfile_path), 'invalid data path : {:}'.format(xfile_path)
+      xdata = torch.load(xfile_path, map_location='cpu')
+      assert isinstance(xdata, dict), 'invalid format of data in {:}'.format(xfile_path)
+
+      hp2archres = OrderedDict()
+      for hp_key, results in xdata.items():
+        hp2archres[hp_key] = ArchResults.create_from_state_dict(results)
+      self.arch2infos_dict[idx] = hp2archres
+
+  def query_info_str_by_arch(self, arch, hp: Text='12'):
+    """ This function is used to query the information of a specific architecture
+        'arch' can be an architecture index or an architecture string
+        When hp=01, the hyper-parameters used to train a model are in 'configs/nas-benchmark/hyper-opts/01E.config'
+        When hp=12, the hyper-parameters used to train a model are in 'configs/nas-benchmark/hyper-opts/12E.config'
+        When hp=90, the hyper-parameters used to train a model are in 'configs/nas-benchmark/hyper-opts/90E.config'
+        The difference between these three configurations are the number of training epochs.
+    """
+    if self.verbose:
+      print('Call query_info_str_by_arch with arch={:} and hp={:}'.format(arch, hp))
+    return self._query_info_str_by_arch(arch, hp, print_information)
+
+  def get_more_info(self, index, dataset: Text, iepoch=None, hp='12', is_random=True):
+    """This function will return the metric for the `index`-th architecture
+       `dataset` indicates the dataset:
+          'cifar10-valid'  : using the proposed train set of CIFAR-10 as the training set
+          'cifar10'        : using the proposed train+valid set of CIFAR-10 as the training set
+          'cifar100'       : using the proposed train set of CIFAR-100 as the training set
+          'ImageNet16-120' : using the proposed train set of ImageNet-16-120 as the training set
+        `iepoch` indicates the index of training epochs from 0 to 11/199.
+          When iepoch=None, it will return the metric for the last training epoch
+          When iepoch=11, it will return the metric for the 11-th training epoch (starting from 0)
+        `hp` indicates different hyper-parameters for training
+          When hp=01, it trains the network with 01 epochs and the LR decayed from 0.1 to 0 within 01 epochs
+          When hp=12, it trains the network with 01 epochs and the LR decayed from 0.1 to 0 within 12 epochs
+          When hp=90, it trains the network with 01 epochs and the LR decayed from 0.1 to 0 within 90 epochs
+        `is_random`
+          When is_random=True, the performance of a random architecture will be returned
+          When is_random=False, the performanceo of all trials will be averaged.
+    """
+    if self.verbose:
+      print('Call the get_more_info function with index={:}, dataset={:}, iepoch={:}, hp={:}, and is_random={:}.'.format(index, dataset, iepoch, hp, is_random))
+    index = self.query_index_by_arch(index)  # To avoid the input is a string or an instance of a arch object
+    if index not in self.arch2infos_dict:
+      raise ValueError('Did not find {:} from arch2infos_dict.'.format(index))
+    archresult = self.arch2infos_dict[index][str(hp)]
+    # if randomly select one trial, select the seed at first
+    if isinstance(is_random, bool) and is_random:
+      seeds = archresult.get_dataset_seeds(dataset)
+      is_random = random.choice(seeds)
+    # collect the training information
+    train_info = archresult.get_metrics(dataset, 'train', iepoch=iepoch, is_random=is_random)
+    total = train_info['iepoch'] + 1
+    xinfo = {'train-loss'    : train_info['loss'],
+             'train-accuracy': train_info['accuracy'],
+             'train-per-time': train_info['all_time'] / total,
+             'train-all-time': train_info['all_time']}
+    # collect the evaluation information
+    if dataset == 'cifar10-valid':
+      valid_info = archresult.get_metrics(dataset, 'x-valid', iepoch=iepoch, is_random=is_random)
+      try:
+        test_info = archresult.get_metrics(dataset, 'ori-test', iepoch=iepoch, is_random=is_random)
+      except:
+        test_info = None
+      valtest_info = None
+    else:
+      try: # collect results on the proposed test set
+        if dataset == 'cifar10':
+          test_info = archresult.get_metrics(dataset, 'ori-test', iepoch=iepoch, is_random=is_random)
+        else:
+          test_info = archresult.get_metrics(dataset, 'x-test', iepoch=iepoch, is_random=is_random)
+      except:
+        test_info = None
+      try: # collect results on the proposed validation set
+        valid_info = archresult.get_metrics(dataset, 'x-valid', iepoch=iepoch, is_random=is_random)
+      except:
+        valid_info = None
+      try:
+        if dataset != 'cifar10':
+          valtest_info = archresult.get_metrics(dataset, 'ori-test', iepoch=iepoch, is_random=is_random)
+        else:
+          valtest_info = None
+      except:
+        valtest_info = None
+    if valid_info is not None:
+      xinfo['valid-loss'] = valid_info['loss']
+      xinfo['valid-accuracy'] = valid_info['accuracy']
+      xinfo['valid-per-time'] = valid_info['all_time'] / total
+      xinfo['valid-all-time'] = valid_info['all_time']
+    if test_info is not None:
+      xinfo['test-loss'] = test_info['loss']
+      xinfo['test-accuracy'] = test_info['accuracy']
+      xinfo['test-per-time'] = test_info['all_time'] / total
+      xinfo['test-all-time'] = test_info['all_time']
+    if valtest_info is not None:
+      xinfo['valtest-loss'] = valtest_info['loss']
+      xinfo['valtest-accuracy'] = valtest_info['accuracy']
+      xinfo['valtest-per-time'] = valtest_info['all_time'] / total
+      xinfo['valtest-all-time'] = valtest_info['all_time']
+    return xinfo
+
+  def show(self, index: int = -1) -> None:
+    """
+    This function will print the information of a specific (or all) architecture(s).
+
+    :param index: If the index < 0: it will loop for all architectures and print their information one by one.
+                  else: it will print the information of the 'index'-th architecture.
+    :return: nothing
+    """
+    self._show(index, print_information)

lib/nats_bench/api_topology.py

@@ -0,0 +1,269 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.07 #
+############################################################################################
+# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size
+############################################################################################
+import os, copy, random, torch, numpy as np
+from pathlib import Path
+from typing import List, Text, Union, Dict, Optional
+from collections import OrderedDict, defaultdict
+
+from .api_utils import ArchResults
+from .api_utils import NASBenchMetaAPI
+from .api_utils import remap_dataset_set_names
+
+
+ALL_BENCHMARK_FILES = ['NAS-Bench-201-v1_0-e61699.pth', 'NAS-Bench-201-v1_1-096897.pth']
+ALL_ARCHIVE_DIRS = ['NAS-Bench-201-v1_1-archive']
+
+
+def print_information(information, extra_info=None, show=False):
+  dataset_names = information.get_dataset_names()
+  strings = [information.arch_str, 'datasets : {:}, extra-info : {:}'.format(dataset_names, extra_info)]
+  def metric2str(loss, acc):
+    return 'loss = {:.3f}, top1 = {:.2f}%'.format(loss, acc)
+
+  for ida, dataset in enumerate(dataset_names):
+    metric = information.get_compute_costs(dataset)
+    flop, param, latency = metric['flops'], metric['params'], metric['latency']
+    str1 = '{:14s} FLOP={:6.2f} M, Params={:.3f} MB, latency={:} ms.'.format(dataset, flop, param, '{:.2f}'.format(latency*1000) if latency is not None and latency > 0 else None)
+    train_info = information.get_metrics(dataset, 'train')
+    if dataset == 'cifar10-valid':
+      valid_info = information.get_metrics(dataset, 'x-valid')
+      str2 = '{:14s} train : [{:}], valid : [{:}]'.format(dataset, metric2str(train_info['loss'], train_info['accuracy']), metric2str(valid_info['loss'], valid_info['accuracy']))
+    elif dataset == 'cifar10':
+      test__info = information.get_metrics(dataset, 'ori-test')
+      str2 = '{:14s} train : [{:}], test  : [{:}]'.format(dataset, metric2str(train_info['loss'], train_info['accuracy']), metric2str(test__info['loss'], test__info['accuracy']))
+    else:
+      valid_info = information.get_metrics(dataset, 'x-valid')
+      test__info = information.get_metrics(dataset, 'x-test')
+      str2 = '{:14s} train : [{:}], valid : [{:}], test : [{:}]'.format(dataset, metric2str(train_info['loss'], train_info['accuracy']), metric2str(valid_info['loss'], valid_info['accuracy']), metric2str(test__info['loss'], test__info['accuracy']))
+    strings += [str1, str2]
+  if show: print('\n'.join(strings))
+  return strings
+
+
+"""
+This is the class for the API of topology search space in NATS-Bench.
+"""
+class NATStopology(NASBenchMetaAPI):
+
+  """ The initialization function that takes the dataset file path (or a dict loaded from that path) as input. """
+  def __init__(self, file_path_or_dict: Optional[Union[Text, Dict]]=None,
+               verbose: bool=True):
+    self.filename = None
+    self._search_space_name = 'topology'
+    self.reset_time()
+    if file_path_or_dict is None:
+      file_path_or_dict = os.path.join(os.environ['TORCH_HOME'], ALL_BENCHMARK_FILES[-1])
+      print ('Try to use the default NATS-Bench (topology) path from {:}.'.format(file_path_or_dict))
+    if isinstance(file_path_or_dict, str) or isinstance(file_path_or_dict, Path):
+      file_path_or_dict = str(file_path_or_dict)
+      if verbose: print('try to create the NATS-Bench (topology) api from {:}'.format(file_path_or_dict))
+      assert os.path.isfile(file_path_or_dict), 'invalid path : {:}'.format(file_path_or_dict)
+      self.filename = Path(file_path_or_dict).name
+      file_path_or_dict = torch.load(file_path_or_dict, map_location='cpu')
+    elif isinstance(file_path_or_dict, dict):
+      file_path_or_dict = copy.deepcopy(file_path_or_dict)
+    else: raise ValueError('invalid type : {:} not in [str, dict]'.format(type(file_path_or_dict)))
+    assert isinstance(file_path_or_dict, dict), 'It should be a dict instead of {:}'.format(type(file_path_or_dict))
+    self.verbose = verbose # [TODO] a flag indicating whether to print more logs
+    keys = ('meta_archs', 'arch2infos', 'evaluated_indexes')
+    for key in keys: assert key in file_path_or_dict, 'Can not find key[{:}] in the dict'.format(key)
+    self.meta_archs = copy.deepcopy( file_path_or_dict['meta_archs'] )
+    # This is a dict mapping each architecture to a dict, where the key is #epochs and the value is ArchResults
+    self.arch2infos_dict = OrderedDict()
+    self._avaliable_hps = set(['12', '200'])
+    for xkey in sorted(list(file_path_or_dict['arch2infos'].keys())):
+      all_info = file_path_or_dict['arch2infos'][xkey]
+      hp2archres = OrderedDict()
+      # self.arch2infos_less[xkey] = ArchResults.create_from_state_dict( all_info['less'] )
+      # self.arch2infos_full[xkey] = ArchResults.create_from_state_dict( all_info['full'] )
+      hp2archres['12'] = ArchResults.create_from_state_dict(all_info['less'])
+      hp2archres['200'] = ArchResults.create_from_state_dict(all_info['full'])
+      self.arch2infos_dict[xkey] = hp2archres
+    self.evaluated_indexes = sorted(list(file_path_or_dict['evaluated_indexes']))
+    self.archstr2index = {}
+    for idx, arch in enumerate(self.meta_archs):
+      assert arch not in self.archstr2index, 'This [{:}]-th arch {:} already in the dict ({:}).'.format(idx, arch, self.archstr2index[arch])
+      self.archstr2index[ arch ] = idx
+
+  def reload(self, archive_root: Text = None, index: int = None):
+    """Overwrite all information of the 'index'-th architecture in the search space.
+         It will load its data from 'archive_root'.
+    """
+    if archive_root is None:
+      archive_root = os.path.join(os.environ['TORCH_HOME'], ALL_ARCHIVE_DIRS[-1])
+    assert os.path.isdir(archive_root), 'invalid directory : {:}'.format(archive_root)
+    if index is None:
+      indexes = list(range(len(self)))
+    else:
+      indexes = [index]
+    for idx in indexes:
+      assert 0 <= idx < len(self.meta_archs), 'invalid index of {:}'.format(idx)
+      xfile_path = os.path.join(archive_root, '{:06d}-FULL.pth'.format(idx))
+      assert os.path.isfile(xfile_path), 'invalid data path : {:}'.format(xfile_path)
+      xdata = torch.load(xfile_path, map_location='cpu')
+      assert isinstance(xdata, dict) and 'full' in xdata and 'less' in xdata, 'invalid format of data in {:}'.format(xfile_path)
+      hp2archres = OrderedDict()
+      hp2archres['12'] = ArchResults.create_from_state_dict(xdata['less'])
+      hp2archres['200'] = ArchResults.create_from_state_dict(xdata['full'])
+      self.arch2infos_dict[idx] = hp2archres
+
+  def query_info_str_by_arch(self, arch, hp: Text='12'):
+    """ This function is used to query the information of a specific architecture
+        'arch' can be an architecture index or an architecture string
+        When hp=12, the hyper-parameters used to train a model are in 'configs/nas-benchmark/hyper-opts/12E.config'
+        When hp=200, the hyper-parameters used to train a model are in 'configs/nas-benchmark/hyper-opts/200E.config'
+        The difference between these three configurations are the number of training epochs.
+    """
+    if self.verbose:
+      print('Call query_info_str_by_arch with arch={:} and hp={:}'.format(arch, hp))
+    return self._query_info_str_by_arch(arch, hp, print_information)
+
+  # obtain the metric for the `index`-th architecture
+  # `dataset` indicates the dataset:
+  #   'cifar10-valid'  : using the proposed train set of CIFAR-10 as the training set
+  #   'cifar10'        : using the proposed train+valid set of CIFAR-10 as the training set
+  #   'cifar100'       : using the proposed train set of CIFAR-100 as the training set
+  #   'ImageNet16-120' : using the proposed train set of ImageNet-16-120 as the training set
+  # `iepoch` indicates the index of training epochs from 0 to 11/199.
+  #   When iepoch=None, it will return the metric for the last training epoch
+  #   When iepoch=11, it will return the metric for the 11-th training epoch (starting from 0)
+  # `use_12epochs_result` indicates different hyper-parameters for training
+  #   When use_12epochs_result=True, it trains the network with 12 epochs and the LR decayed from 0.1 to 0 within 12 epochs
+  #   When use_12epochs_result=False, it trains the network with 200 epochs and the LR decayed from 0.1 to 0 within 200 epochs
+  # `is_random`
+  #   When is_random=True, the performance of a random architecture will be returned
+  #   When is_random=False, the performanceo of all trials will be averaged.
+  def get_more_info(self, index, dataset, iepoch=None, hp='12', is_random=True):
+    if self.verbose:
+      print('Call the get_more_info function with index={:}, dataset={:}, iepoch={:}, hp={:}, and is_random={:}.'.format(index, dataset, iepoch, hp, is_random))
+    index = self.query_index_by_arch(index)  # To avoid the input is a string or an instance of a arch object
+    if index not in self.arch2infos_dict:
+      raise ValueError('Did not find {:} from arch2infos_dict.'.format(index))
+    archresult = self.arch2infos_dict[index][str(hp)]
+    # if randomly select one trial, select the seed at first
+    if isinstance(is_random, bool) and is_random:
+      seeds = archresult.get_dataset_seeds(dataset)
+      is_random = random.choice(seeds)
+    # collect the training information
+    train_info = archresult.get_metrics(dataset, 'train', iepoch=iepoch, is_random=is_random)
+    total = train_info['iepoch'] + 1
+    xinfo = {'train-loss'    : train_info['loss'],
+             'train-accuracy': train_info['accuracy'],
+             'train-per-time': train_info['all_time'] / total if train_info['all_time'] is not None else None,
+             'train-all-time': train_info['all_time']}
+    # collect the evaluation information
+    if dataset == 'cifar10-valid':
+      valid_info = archresult.get_metrics(dataset, 'x-valid', iepoch=iepoch, is_random=is_random)
+      try:
+        test_info = archresult.get_metrics(dataset, 'ori-test', iepoch=iepoch, is_random=is_random)
+      except:
+        test_info = None
+      valtest_info = None
+    else:
+      try: # collect results on the proposed test set
+        if dataset == 'cifar10':
+          test_info = archresult.get_metrics(dataset, 'ori-test', iepoch=iepoch, is_random=is_random)
+        else:
+          test_info = archresult.get_metrics(dataset, 'x-test', iepoch=iepoch, is_random=is_random)
+      except:
+        test_info = None
+      try: # collect results on the proposed validation set
+        valid_info = archresult.get_metrics(dataset, 'x-valid', iepoch=iepoch, is_random=is_random)
+      except:
+        valid_info = None
+      try:
+        if dataset != 'cifar10':
+          valtest_info = archresult.get_metrics(dataset, 'ori-test', iepoch=iepoch, is_random=is_random)
+        else:
+          valtest_info = None
+      except:
+        valtest_info = None
+    if valid_info is not None:
+      xinfo['valid-loss'] = valid_info['loss']
+      xinfo['valid-accuracy'] = valid_info['accuracy']
+      xinfo['valid-per-time'] = valid_info['all_time'] / total if valid_info['all_time'] is not None else None
+      xinfo['valid-all-time'] = valid_info['all_time']
+    if test_info is not None:
+      xinfo['test-loss'] = test_info['loss']
+      xinfo['test-accuracy'] = test_info['accuracy']
+      xinfo['test-per-time'] = test_info['all_time'] / total if test_info['all_time'] is not None else None
+      xinfo['test-all-time'] = test_info['all_time']
+    if valtest_info is not None:
+      xinfo['valtest-loss'] = valtest_info['loss']
+      xinfo['valtest-accuracy'] = valtest_info['accuracy']
+      xinfo['valtest-per-time'] = valtest_info['all_time'] / total if valtest_info['all_time'] is not None else None
+      xinfo['valtest-all-time'] = valtest_info['all_time']
+    return xinfo
+
+  def show(self, index: int = -1) -> None:
+    """This function will print the information of a specific (or all) architecture(s)."""
+    self._show(index, print_information)
+
+  @staticmethod
+  def str2lists(arch_str: Text) -> List[tuple]:
+    """
+    This function shows how to read the string-based architecture encoding.
+      It is the same as the `str2structure` func in `AutoDL-Projects/lib/models/cell_searchs/genotypes.py`
+
+    :param
+      arch_str: the input is a string indicates the architecture topology, such as
+                    |nor_conv_1x1~0|+|none~0|none~1|+|none~0|none~1|skip_connect~2|
+    :return: a list of tuple, contains multiple (op, input_node_index) pairs.
+
+    :usage
+      arch = api.str2lists( '|nor_conv_1x1~0|+|none~0|none~1|+|none~0|none~1|skip_connect~2|' )
+      print ('there are {:} nodes in this arch'.format(len(arch)+1)) # arch is a list
+      for i, node in enumerate(arch):
+        print('the {:}-th node is the sum of these {:} nodes with op: {:}'.format(i+1, len(node), node))
+    """
+    node_strs = arch_str.split('+')
+    genotypes = []
+    for i, node_str in enumerate(node_strs):
+      inputs = list(filter(lambda x: x != '', node_str.split('|')))
+      for xinput in inputs: assert len(xinput.split('~')) == 2, 'invalid input length : {:}'.format(xinput)
+      inputs = ( xi.split('~') for xi in inputs )
+      input_infos = tuple( (op, int(IDX)) for (op, IDX) in inputs)
+      genotypes.append( input_infos )
+    return genotypes
+
+  @staticmethod
+  def str2matrix(arch_str: Text,
+                 search_space: List[Text] = ['none', 'skip_connect', 'nor_conv_1x1', 'nor_conv_3x3', 'avg_pool_3x3']) -> np.ndarray:
+    """
+    This func shows how to convert the string-based architecture encoding to the encoding strategy in NAS-Bench-101.
+
+    :param
+      arch_str: the input is a string indicates the architecture topology, such as
+                    |nor_conv_1x1~0|+|none~0|none~1|+|none~0|none~1|skip_connect~2|
+      search_space: a list of operation string, the default list is the topology search space for NATS-BENCH.
+        the default value should be be consistent with this line https://github.com/D-X-Y/AutoDL-Projects/blob/master/lib/models/cell_operations.py#L24
+    :return
+      the numpy matrix (2-D np.ndarray) representing the DAG of this architecture topology
+    :usage
+      matrix = api.str2matrix( '|nor_conv_1x1~0|+|none~0|none~1|+|none~0|none~1|skip_connect~2|' )
+      This matrix is 4-by-4 matrix representing a cell with 4 nodes (only the lower left triangle is useful).
+         [ [0, 0, 0, 0],  # the first line represents the input (0-th) node
+           [2, 0, 0, 0],  # the second line represents the 1-st node, is calculated by 2-th-op( 0-th-node )
+           [0, 0, 0, 0],  # the third line represents the 2-nd node, is calculated by 0-th-op( 0-th-node ) + 0-th-op( 1-th-node )
+           [0, 0, 1, 0] ] # the fourth line represents the 3-rd node, is calculated by 0-th-op( 0-th-node ) + 0-th-op( 1-th-node ) + 1-th-op( 2-th-node )
+      In the topology search space in NATS-BENCH, 0-th-op is 'none', 1-th-op is 'skip_connect',
+         2-th-op is 'nor_conv_1x1', 3-th-op is 'nor_conv_3x3', 4-th-op is 'avg_pool_3x3'.
+    :(NOTE)
+      If a node has two input-edges from the same node, this function does not work. One edge will be overlapped.
+    """
+    node_strs = arch_str.split('+')
+    num_nodes = len(node_strs) + 1
+    matrix = np.zeros((num_nodes, num_nodes))
+    for i, node_str in enumerate(node_strs):
+      inputs = list(filter(lambda x: x != '', node_str.split('|')))
+      for xinput in inputs: assert len(xinput.split('~')) == 2, 'invalid input length : {:}'.format(xinput)
+      for xi in inputs:
+        op, idx = xi.split('~')
+        if op not in search_space: raise ValueError('this op ({:}) is not in {:}'.format(op, search_space))
+        op_idx, node_idx = search_space.index(op), int(idx)
+        matrix[i+1, node_idx] = op_idx
+    return matrix
+

lib/nats_bench/api_utils.py

@@ -0,0 +1,754 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.07 #
+############################################################################################
+# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size
+############################################################################################
+# In this Python file, we define NASBenchMetaAPI, the abstract class for benchmark APIs.
+# We also define the class ArchResults, which contains all information of a single architecture trained by one kind of hyper-parameters on three datasets.
+# We also define the class ResultsCount, which contains all information of a single trial for a single architecture.
+############################################################################################
+# History:
+# [2020.07.31] The first version, where most content reused nas_201_api/api_utils.py
+#
+import os, abc, copy, random, torch, numpy as np
+from pathlib import Path
+from typing import List, Text, Union, Dict, Optional
+from collections import OrderedDict, defaultdict
+
+
+def remap_dataset_set_names(dataset, metric_on_set, verbose=False):
+  """re-map the metric_on_set to internal keys"""
+  if verbose:
+    print('Call internal function _remap_dataset_set_names with dataset={:} and metric_on_set={:}'.format(dataset, metric_on_set))
+  if dataset == 'cifar10' and metric_on_set == 'valid':
+    dataset, metric_on_set = 'cifar10-valid', 'x-valid'
+  elif dataset == 'cifar10' and metric_on_set == 'test':
+    dataset, metric_on_set = 'cifar10', 'ori-test'
+  elif dataset == 'cifar10' and metric_on_set == 'train':
+    dataset, metric_on_set = 'cifar10', 'train'
+  elif (dataset == 'cifar100' or dataset == 'ImageNet16-120') and metric_on_set == 'valid':
+    metric_on_set = 'x-valid'
+  elif (dataset == 'cifar100' or dataset == 'ImageNet16-120') and metric_on_set == 'test':
+    metric_on_set = 'x-test'
+  if verbose:
+    print('  return dataset={:} and metric_on_set={:}'.format(dataset, metric_on_set))
+  return dataset, metric_on_set
+
+
+class NASBenchMetaAPI(metaclass=abc.ABCMeta):
+
+  @abc.abstractmethod
+  def __init__(self, file_path_or_dict: Optional[Union[Text, Dict]]=None, verbose: bool=True):
+    """The initialization function that takes the dataset file path (or a dict loaded from that path) as input."""
+
+  def __getitem__(self, index: int):
+    return copy.deepcopy(self.meta_archs[index])
+
+  def arch(self, index: int):
+    """Return the topology structure of the `index`-th architecture."""
+    if self.verbose:
+      print('Call the arch function with index={:}'.format(index))
+    assert 0 <= index < len(self.meta_archs), 'invalid index : {:} vs. {:}.'.format(index, len(self.meta_archs))
+    return copy.deepcopy(self.meta_archs[index])
+
+  def __len__(self):
+    return len(self.meta_archs)
+
+  def __repr__(self):
+    return ('{name}({num}/{total} architectures, file={filename})'.format(name=self.__class__.__name__, num=len(self.evaluated_indexes), total=len(self.meta_archs), filename=self.filename))
+
+  @property
+  def avaliable_hps(self):
+    return list(copy.deepcopy(self._avaliable_hps))
+
+  @property
+  def used_time(self):
+    return self._used_time
+  
+  @property
+  def search_space_name(self):
+    return self._search_space_name
+
+  def reset_time(self):
+    self._used_time = 0
+
+  def simulate_train_eval(self, arch, dataset, iepoch=None, hp='12', account_time=True):
+    index = self.query_index_by_arch(arch)
+    all_names = ('cifar10', 'cifar100', 'ImageNet16-120')
+    assert dataset in all_names, 'Invalid dataset name : {:} vs {:}'.format(dataset, all_names)
+    if dataset == 'cifar10':
+      info = self.get_more_info(index, 'cifar10-valid', iepoch=iepoch, hp=hp, is_random=True)
+    else:
+      info = self.get_more_info(index, dataset, iepoch=iepoch, hp=hp, is_random=True)
+    valid_acc, time_cost = info['valid-accuracy'], info['train-all-time'] + info['valid-per-time']
+    latency = self.get_latency(index, dataset)
+    if account_time:
+      self._used_time += time_cost
+    return valid_acc, latency, time_cost, self._used_time
+
+  def random(self):
+    """Return a random index of all architectures."""
+    return random.randint(0, len(self.meta_archs)-1)
+
+  def query_index_by_arch(self, arch):
+    """ This function is used to query the index of an architecture in the search space.
+        In the topology search space, the input arch can be an architecture string such as '|nor_conv_3x3~0|+|nor_conv_3x3~0|avg_pool_3x3~1|+|skip_connect~0|nor_conv_3x3~1|skip_connect~2|';
+          or an instance that has the 'tostr' function that can generate the architecture string;
+          or it is directly an architecture index, in this case, we will check whether it is valid or not.
+        This function will return the index.
+        If return -1, it means this architecture is not in the search space.
+        Otherwise, it will return an int in [0, the-number-of-candidates-in-the-search-space).
+    """
+    if self.verbose:
+      print('Call query_index_by_arch with arch={:}'.format(arch))
+    if isinstance(arch, int):
+      if 0 <= arch < len(self):
+        return arch
+      else:
+        raise ValueError('Invalid architecture index {:} vs [{:}, {:}].'.format(arch, 0, len(self)))
+    elif isinstance(arch, str):
+      if arch in self.archstr2index: arch_index = self.archstr2index[ arch ]
+      else                         : arch_index = -1
+    elif hasattr(arch, 'tostr'):
+      if arch.tostr() in self.archstr2index: arch_index = self.archstr2index[ arch.tostr() ]
+      else                                 : arch_index = -1
+    else: arch_index = -1
+    return arch_index
+
+  def query_by_arch(self, arch, hp):
+    # This is to make the current version be compatible with the old version.
+    return self.query_info_str_by_arch(arch, hp)
+
+  @abc.abstractmethod
+  def reload(self, archive_root: Text = None, index: int = None):
+    """Overwrite all information of the 'index'-th architecture in the search space, where the data will be loaded from 'archive_root'.
+       If index is None, overwrite all ckps.
+    """
+
+  def clear_params(self, index: int, hp: Optional[Text]=None):
+    """Remove the architecture's weights to save memory.
+    :arg
+      index: the index of the target architecture
+      hp: a flag to controll how to clear the parameters.
+        -- None: clear all the weights in '01'/'12'/'90', which indicates the number of training epochs.
+        -- '01' or '12' or '90': clear all the weights in arch2infos_dict[index][hp].
+    """
+    if self.verbose:
+      print('Call clear_params with index={:} and hp={:}'.format(index, hp))
+    if hp is None:
+      for key, result in self.arch2infos_dict[index].items():
+        result.clear_params()
+    else:
+      if str(hp) not in self.arch2infos_dict[index]:
+        raise ValueError('The {:}-th architecture only has hyper-parameters of {:} instead of {:}.'.format(index, list(self.arch2infos_dict[index].keys()), hp))
+      self.arch2infos_dict[index][str(hp)].clear_params()
+
+  @abc.abstractmethod
+  def query_info_str_by_arch(self, arch, hp: Text='12'):
+    """This function is used to query the information of a specific architecture."""
+
+  def _query_info_str_by_arch(self, arch, hp: Text='12', print_information=None):
+    arch_index = self.query_index_by_arch(arch)
+    if arch_index in self.arch2infos_dict:
+      if hp not in self.arch2infos_dict[arch_index]:
+        raise ValueError('The {:}-th architecture only has hyper-parameters of {:} instead of {:}.'.format(index, list(self.arch2infos_dict[arch_index].keys()), hp))
+      info = self.arch2infos_dict[arch_index][hp]
+      strings = print_information(info, 'arch-index={:}'.format(arch_index))
+      return '\n'.join(strings)
+    else:
+      print ('Find this arch-index : {:}, but this arch is not evaluated.'.format(arch_index))
+      return None
+
+  def query_meta_info_by_index(self, arch_index, hp: Text = '12'):
+    """Return the ArchResults for the 'arch_index'-th architecture. This function is similar to query_by_index."""
+    if self.verbose:
+      print('Call query_meta_info_by_index with arch_index={:}, hp={:}'.format(arch_index, hp))
+    if arch_index in self.arch2infos_dict:
+      if hp not in self.arch2infos_dict[arch_index]:
+        raise ValueError('The {:}-th architecture only has hyper-parameters of {:} instead of {:}.'.format(arch_index, list(self.arch2infos_dict[arch_index].keys()), hp))
+      info = self.arch2infos_dict[arch_index][hp]
+    else:
+      raise ValueError('arch_index [{:}] does not in arch2infos'.format(arch_index))
+    return copy.deepcopy(info)
+
+  def query_by_index(self, arch_index: int, dataname: Union[None, Text] = None, hp: Text = '12'):
+    """ This 'query_by_index' function is used to query information with the training of 01 epochs, 12 epochs, 90 epochs, or 200 epochs.
+        ------
+        If hp=01, we train the model by 01 epochs (see config in configs/nas-benchmark/hyper-opts/01E.config)
+        If hp=12, we train the model by 01 epochs (see config in configs/nas-benchmark/hyper-opts/12E.config)
+        If hp=90, we train the model by 01 epochs (see config in configs/nas-benchmark/hyper-opts/90E.config)
+        If hp=200, we train the model by 01 epochs (see config in configs/nas-benchmark/hyper-opts/200E.config)
+        ------
+        If dataname is None, return the ArchResults
+          else, return a dict with all trials on that dataset (the key is the seed)
+        Options are 'cifar10-valid', 'cifar10', 'cifar100', 'ImageNet16-120'.
+        -- cifar10-valid : training the model on the CIFAR-10 training set.
+        -- cifar10 : training the model on the CIFAR-10 training + validation set.
+        -- cifar100 : training the model on the CIFAR-100 training set.
+        -- ImageNet16-120 : training the model on the ImageNet16-120 training set.
+    """
+    if self.verbose:
+      print('Call query_by_index with arch_index={:}, dataname={:}, hp={:}'.format(arch_index, dataname, hp))
+    info = self.query_meta_info_by_index(arch_index, hp)
+    if dataname is None: return info
+    else:
+      if dataname not in info.get_dataset_names():
+        raise ValueError('invalid dataset-name : {:} vs. {:}'.format(dataname, info.get_dataset_names()))
+      return info.query(dataname)
+
+  def find_best(self, dataset, metric_on_set, FLOP_max=None, Param_max=None, hp: Text = '12'):
+    """Find the architecture with the highest accuracy based on some constraints."""
+    if self.verbose:
+      print('Call find_best with dataset={:}, metric_on_set={:}, hp={:} | with #FLOPs < {:} and #Params < {:}'.format(dataset, metric_on_set, hp, FLOP_max, Param_max))
+    dataset, metric_on_set = remap_dataset_set_names(dataset, metric_on_set, self.verbose)
+    best_index, highest_accuracy = -1, None
+    for i, arch_index in enumerate(self.evaluated_indexes):
+      arch_info = self.arch2infos_dict[arch_index][hp]
+      info = arch_info.get_compute_costs(dataset)  # the information of costs
+      flop, param, latency = info['flops'], info['params'], info['latency']
+      if FLOP_max  is not None and flop  > FLOP_max : continue
+      if Param_max is not None and param > Param_max: continue
+      xinfo = arch_info.get_metrics(dataset, metric_on_set)  # the information of loss and accuracy
+      loss, accuracy = xinfo['loss'], xinfo['accuracy']
+      if best_index == -1:
+        best_index, highest_accuracy = arch_index, accuracy
+      elif highest_accuracy < accuracy:
+        best_index, highest_accuracy = arch_index, accuracy
+    if self.verbose:
+      print('  the best architecture : [{:}] {:} with accuracy={:.3f}%'.format(best_index, self.arch(best_index), highest_accuracy))
+    return best_index, highest_accuracy
+
+  def get_net_param(self, index, dataset, seed: Optional[int], hp: Text = '12'):
+    """
+      This function is used to obtain the trained weights of the `index`-th architecture on `dataset` with the seed of `seed`
+      Args [seed]:
+        -- None : return a dict containing the trained weights of all trials, where each key is a seed and its corresponding value is the weights.
+        -- a interger : return the weights of a specific trial, whose seed is this interger.
+      Args [hp]:
+        -- 01 : train the model by 01 epochs
+        -- 12 : train the model by 12 epochs
+        -- 90 : train the model by 90 epochs
+        -- 200 : train the model by 200 epochs
+    """
+    if self.verbose:
+      print('Call the get_net_param function with index={:}, dataset={:}, seed={:}, hp={:}'.format(index, dataset, seed, hp))
+    info = self.query_meta_info_by_index(index, hp)
+    return info.get_net_param(dataset, seed)
+
+  def get_net_config(self, index: int, dataset: Text):
+    """
+      This function is used to obtain the configuration for the `index`-th architecture on `dataset`.
+      Args [dataset] (4 possible options):
+        -- cifar10-valid : training the model on the CIFAR-10 training set.
+        -- cifar10 : training the model on the CIFAR-10 training + validation set.
+        -- cifar100 : training the model on the CIFAR-100 training set.
+        -- ImageNet16-120 : training the model on the ImageNet16-120 training set.
+      This function will return a dict.
+      ========= Some examlpes for using this function:
+      config = api.get_net_config(128, 'cifar10')
+    """
+    if self.verbose:
+      print('Call the get_net_config function with index={:}, dataset={:}.'.format(index, dataset))
+    if index in self.arch2infos_dict:
+      info = self.arch2infos_dict[index]
+    else:
+      raise ValueError('The arch_index={:} is not in arch2infos_dict.'.format(arch_index))
+    info = next(iter(info.values()))
+    results = info.query(dataset, None)
+    results = next(iter(results.values()))
+    return results.get_config(None)
+  
+  def get_cost_info(self, index: int, dataset: Text, hp: Text = '12') -> Dict[Text, float]:
+    """To obtain the cost metric for the `index`-th architecture on a dataset."""
+    if self.verbose:
+      print('Call the get_cost_info function with index={:}, dataset={:}, and hp={:}.'.format(index, dataset, hp))
+    info = self.query_meta_info_by_index(index, hp)
+    return info.get_compute_costs(dataset)
+
+  def get_latency(self, index: int, dataset: Text, hp: Text = '12') -> float:
+    """
+    To obtain the latency of the network (by default it will return the latency with the batch size of 256).
+    :param index: the index of the target architecture
+    :param dataset: the dataset name (cifar10-valid, cifar10, cifar100, ImageNet16-120)
+    :return: return a float value in seconds
+    """
+    if self.verbose:
+      print('Call the get_latency function with index={:}, dataset={:}, and hp={:}.'.format(index, dataset, hp))
+    cost_dict = self.get_cost_info(index, dataset, hp)
+    return cost_dict['latency']
+
+  @abc.abstractmethod
+  def show(self, index=-1):
+    """This function will print the information of a specific (or all) architecture(s)."""
+
+  def _show(self, index=-1, print_information=None) -> None:
+    """
+    This function will print the information of a specific (or all) architecture(s).
+
+    :param index: If the index < 0: it will loop for all architectures and print their information one by one.
+                  else: it will print the information of the 'index'-th architecture.
+    :return: nothing
+    """
+    if index < 0: # show all architectures
+      print(self)
+      for i, idx in enumerate(self.evaluated_indexes):
+        print('\n' + '-' * 10 + ' The ({:5d}/{:5d}) {:06d}-th architecture! '.format(i, len(self.evaluated_indexes), idx) + '-'*10)
+        print('arch : {:}'.format(self.meta_archs[idx]))
+        for key, result in self.arch2infos_dict[index].items():
+          strings = print_information(result)
+          print('>' * 40 + ' {:03d} epochs '.format(result.get_total_epoch()) + '>' * 40)
+          print('\n'.join(strings))
+        print('<' * 40 + '------------' + '<' * 40)
+    else:
+      if 0 <= index < len(self.meta_archs):
+        if index not in self.evaluated_indexes: print('The {:}-th architecture has not been evaluated or not saved.'.format(index))
+        else:
+          arch_info = self.arch2infos_dict[index]
+          for key, result in self.arch2infos_dict[index].items():
+            strings = print_information(result)
+            print('>' * 40 + ' {:03d} epochs '.format(result.get_total_epoch()) + '>' * 40)
+            print('\n'.join(strings))
+          print('<' * 40 + '------------' + '<' * 40)
+      else:
+        print('This index ({:}) is out of range (0~{:}).'.format(index, len(self.meta_archs)))
+
+  def statistics(self, dataset: Text, hp: Union[Text, int]) -> Dict[int, int]:
+    """This function will count the number of total trials."""
+    if self.verbose:
+      print('Call the statistics function with dataset={:} and hp={:}.'.format(dataset, hp))
+    valid_datasets = ['cifar10-valid', 'cifar10', 'cifar100', 'ImageNet16-120']
+    if dataset not in valid_datasets:
+      raise ValueError('{:} not in {:}'.format(dataset, valid_datasets))
+    nums, hp = defaultdict(lambda: 0), str(hp)
+    for index in range(len(self)):
+      archInfo = self.arch2infos_dict[index][hp]
+      dataset_seed = archInfo.dataset_seed
+      if dataset not in dataset_seed:
+        nums[0] += 1
+      else:
+        nums[len(dataset_seed[dataset])] += 1
+    return dict(nums)
+
+
+class ArchResults(object):
+
+  def __init__(self, arch_index, arch_str):
+    self.arch_index   = int(arch_index)
+    self.arch_str     = copy.deepcopy(arch_str)
+    self.all_results  = dict()
+    self.dataset_seed = dict()
+    self.clear_net_done = False
+
+  def get_compute_costs(self, dataset):
+    x_seeds = self.dataset_seed[dataset]
+    results = [self.all_results[ (dataset, seed) ] for seed in x_seeds]
+
+    flops     = [result.flop for result in results]
+    params    = [result.params for result in results]
+    latencies = [result.get_latency() for result in results]
+    latencies = [x for x in latencies if x > 0]
+    mean_latency = np.mean(latencies) if len(latencies) > 0 else None
+    time_infos = defaultdict(list)
+    for result in results:
+      time_info = result.get_times()
+      for key, value in time_info.items(): time_infos[key].append( value )
+     
+    info = {'flops'  : np.mean(flops),
+            'params' : np.mean(params),
+            'latency': mean_latency}
+    for key, value in time_infos.items():
+      if len(value) > 0 and value[0] is not None:
+        info[key] = np.mean(value)
+      else: info[key] = None
+    return info
+
+  def get_metrics(self, dataset, setname, iepoch=None, is_random=False):
+    """
+      This `get_metrics` function is used to obtain obtain the loss, accuracy, etc information on a specific dataset.
+      If not specify, each set refer to the proposed split in NAS-Bench-201 paper.
+      If some args return None or raise error, then it is not avaliable.
+      ========================================
+      Args [dataset] (4 possible options):
+        -- cifar10-valid : training the model on the CIFAR-10 training set.
+        -- cifar10 : training the model on the CIFAR-10 training + validation set.
+        -- cifar100 : training the model on the CIFAR-100 training set.
+        -- ImageNet16-120 : training the model on the ImageNet16-120 training set.
+      Args [setname] (each dataset has different setnames):
+        -- When dataset = cifar10-valid, you can use 'train', 'x-valid', 'ori-test'
+        ------ 'train' : the metric on the training set.
+        ------ 'x-valid' : the metric on the validation set.
+        ------ 'ori-test' : the metric on the test set.
+        -- When dataset = cifar10, you can use 'train', 'ori-test'.
+        ------ 'train' : the metric on the training + validation set.
+        ------ 'ori-test' : the metric on the test set.
+        -- When dataset = cifar100 or ImageNet16-120, you can use 'train', 'ori-test', 'x-valid', 'x-test'
+        ------ 'train' : the metric on the training set.
+        ------ 'x-valid' : the metric on the validation set.
+        ------ 'x-test' : the metric on the test set.
+        ------ 'ori-test' : the metric on the validation + test set.
+      Args [iepoch] (None or an integer in [0, the-number-of-total-training-epochs)
+        ------ None : return the metric after the last training epoch.
+        ------ an integer i : return the metric after the i-th training epoch.
+      Args [is_random]:
+        ------ True : return the metric of a randomly selected trial.
+        ------ False : return the averaged metric of all avaliable trials.
+        ------ an integer indicating the 'seed' value : return the metric of a specific trial (whose random seed is 'is_random').
+    """
+    x_seeds = self.dataset_seed[dataset]
+    results = [self.all_results[ (dataset, seed) ] for seed in x_seeds]
+    infos   = defaultdict(list)
+    for result in results:
+      if setname == 'train':
+        info = result.get_train(iepoch)
+      else:
+        info = result.get_eval(setname, iepoch)
+      for key, value in info.items(): infos[key].append( value )
+    return_info = dict()
+    if isinstance(is_random, bool) and is_random: # randomly select one
+      index = random.randint(0, len(results)-1)
+      for key, value in infos.items(): return_info[key] = value[index]
+    elif isinstance(is_random, bool) and not is_random: # average
+      for key, value in infos.items():
+        if len(value) > 0 and value[0] is not None:
+          return_info[key] = np.mean(value)
+        else: return_info[key] = None
+    elif isinstance(is_random, int): # specify the seed
+      if is_random not in x_seeds: raise ValueError('can not find random seed ({:}) from {:}'.format(is_random, x_seeds))
+      index = x_seeds.index(is_random)
+      for key, value in infos.items(): return_info[key] = value[index]
+    else:
+      raise ValueError('invalid value for is_random: {:}'.format(is_random))
+    return return_info
+
+  def show(self, is_print=False):
+    return print_information(self, None, is_print)
+
+  def get_dataset_names(self):
+    return list(self.dataset_seed.keys())
+
+  def get_dataset_seeds(self, dataset):
+    return copy.deepcopy( self.dataset_seed[dataset] )
+
+  def get_net_param(self, dataset: Text, seed: Union[None, int] =None):
+    """
+    This function will return the trained network's weights on the 'dataset'.
+    :arg
+      dataset: one of 'cifar10-valid', 'cifar10', 'cifar100', and 'ImageNet16-120'.
+      seed: an integer indicates the seed value or None that indicates returing all trials.
+    """
+    if seed is None:
+      x_seeds = self.dataset_seed[dataset]
+      return {seed: self.all_results[(dataset, seed)].get_net_param() for seed in x_seeds}
+    else:
+      xkey = (dataset, seed)
+      if xkey in self.all_results:
+        return self.all_results[xkey].get_net_param()
+      else:
+        raise ValueError('key={:} not in {:}'.format(xkey, list(self.all_results.keys())))
+
+  def reset_latency(self, dataset: Text, seed: Union[None, Text], latency: float) -> None:
+    """This function is used to reset the latency in all corresponding ResultsCount(s)."""
+    if seed is None:
+      for seed in self.dataset_seed[dataset]:
+        self.all_results[(dataset, seed)].update_latency([latency])
+    else:
+      self.all_results[(dataset, seed)].update_latency([latency])
+
+  def reset_pseudo_train_times(self, dataset: Text, seed: Union[None, Text], estimated_per_epoch_time: float) -> None:
+    """This function is used to reset the train-times in all corresponding ResultsCount(s)."""
+    if seed is None:
+      for seed in self.dataset_seed[dataset]:
+        self.all_results[(dataset, seed)].reset_pseudo_train_times(estimated_per_epoch_time)
+    else:
+      self.all_results[(dataset, seed)].reset_pseudo_train_times(estimated_per_epoch_time)
+
+  def reset_pseudo_eval_times(self, dataset: Text, seed: Union[None, Text], eval_name: Text, estimated_per_epoch_time: float) -> None:
+    """This function is used to reset the eval-times in all corresponding ResultsCount(s)."""
+    if seed is None:
+      for seed in self.dataset_seed[dataset]:
+        self.all_results[(dataset, seed)].reset_pseudo_eval_times(eval_name, estimated_per_epoch_time)
+    else:
+      self.all_results[(dataset, seed)].reset_pseudo_eval_times(eval_name, estimated_per_epoch_time)
+
+  def get_latency(self, dataset: Text) -> float:
+    """Get the latency of a model on the target dataset. [Timestamp: 2020.03.09]"""
+    latencies = []
+    for seed in self.dataset_seed[dataset]:
+      latency = self.all_results[(dataset, seed)].get_latency()
+      if not isinstance(latency, float) or latency <= 0:
+        raise ValueError('invalid latency of {:} with seed={:} : {:}'.format(dataset, seed, latency))
+      latencies.append(latency)
+    return sum(latencies) / len(latencies)
+
+  def get_total_epoch(self, dataset=None):
+    """Return the total number of training epochs."""
+    if dataset is None:
+      epochss = []
+      for xdata, x_seeds in self.dataset_seed.items():
+        epochss += [self.all_results[(xdata, seed)].get_total_epoch() for seed in x_seeds]
+    elif isinstance(dataset, str):
+      x_seeds = self.dataset_seed[dataset]
+      epochss = [self.all_results[(dataset, seed)].get_total_epoch() for seed in x_seeds]
+    else:
+      raise ValueError('invalid dataset={:}'.format(dataset))
+    if len(set(epochss)) > 1: raise ValueError('Each trial mush have the same number of training epochs : {:}'.format(epochss))
+    return epochss[-1]
+
+  def query(self, dataset, seed=None):
+    """Return the ResultsCount object (containing all information of a single trial) for 'dataset' and 'seed'"""
+    if seed is None:
+      x_seeds = self.dataset_seed[dataset]
+      return {seed: self.all_results[(dataset, seed)] for seed in x_seeds}
+    else:
+      return self.all_results[(dataset, seed)]
+
+  def arch_idx_str(self):
+    return '{:06d}'.format(self.arch_index)
+
+  def update(self, dataset_name, seed, result):
+    if dataset_name not in self.dataset_seed:
+      self.dataset_seed[dataset_name] = []
+    assert seed not in self.dataset_seed[dataset_name], '{:}-th arch alreadly has this seed ({:}) on {:}'.format(self.arch_index, seed, dataset_name)
+    self.dataset_seed[ dataset_name ].append( seed )
+    self.dataset_seed[ dataset_name ] = sorted( self.dataset_seed[ dataset_name ] )
+    assert (dataset_name, seed) not in self.all_results
+    self.all_results[ (dataset_name, seed) ] = result
+    self.clear_net_done = False
+
+  def state_dict(self):
+    state_dict = dict()
+    for key, value in self.__dict__.items():
+      if key == 'all_results': # contain the class of ResultsCount
+        xvalue = dict()
+        assert isinstance(value, dict), 'invalid type of value for {:} : {:}'.format(key, type(value))
+        for _k, _v in value.items():
+          assert isinstance(_v, ResultsCount), 'invalid type of value for {:}/{:} : {:}'.format(key, _k, type(_v))
+          xvalue[_k] = _v.state_dict()
+      else:
+        xvalue = value
+      state_dict[key] = xvalue
+    return state_dict
+
+  def load_state_dict(self, state_dict):
+    new_state_dict = dict()
+    for key, value in state_dict.items():
+      if key == 'all_results': # to convert to the class of ResultsCount
+        xvalue = dict()
+        assert isinstance(value, dict), 'invalid type of value for {:} : {:}'.format(key, type(value))
+        for _k, _v in value.items():
+          xvalue[_k] = ResultsCount.create_from_state_dict(_v)
+      else: xvalue = value
+      new_state_dict[key] = xvalue
+    self.__dict__.update(new_state_dict)
+
+  @staticmethod
+  def create_from_state_dict(state_dict_or_file):
+    x = ArchResults(-1, -1)
+    if isinstance(state_dict_or_file, str): # a file path
+      state_dict = torch.load(state_dict_or_file, map_location='cpu')
+    elif isinstance(state_dict_or_file, dict):
+      state_dict = state_dict_or_file
+    else:
+      raise ValueError('invalid type of state_dict_or_file : {:}'.format(type(state_dict_or_file)))
+    x.load_state_dict(state_dict)
+    return x
+
+  # This function is used to clear the weights saved in each 'result'
+  # This can help reduce the memory footprint.
+  def clear_params(self):
+    for key, result in self.all_results.items():
+      del result.net_state_dict
+      result.net_state_dict = None
+    self.clear_net_done = True
+
+  def debug_test(self):
+    """This function is used for me to debug and test, which will call most methods."""
+    all_dataset = ['cifar10-valid', 'cifar10', 'cifar100', 'ImageNet16-120']
+    for dataset in all_dataset:
+      print('---->>>> {:}'.format(dataset))
+      print('The latency on {:} is {:} s'.format(dataset, self.get_latency(dataset)))
+      for seed in self.dataset_seed[dataset]:
+        result = self.all_results[(dataset, seed)]
+        print('  ==>> result = {:}'.format(result))
+        print('  ==>> cost = {:}'.format(result.get_times()))
+
+  def __repr__(self):
+    return ('{name}(arch-index={index}, arch={arch}, {num} runs, clear={clear})'.format(name=self.__class__.__name__, index=self.arch_index, arch=self.arch_str, num=len(self.all_results), clear=self.clear_net_done))
+
+
+"""
+This class (ResultsCount) is used to save the information of one trial for a single architecture.
+I did not write much comment for this class, because it is the lowest-level class in NAS-Bench-201 API, which will be rarely called.
+If you have any question regarding this class, please open an issue or email me.
+"""
+class ResultsCount(object):
+
+  def __init__(self, name, state_dict, train_accs, train_losses, params, flop, arch_config, seed, epochs, latency):
+    self.name           = name
+    self.net_state_dict = state_dict
+    self.train_acc1es = copy.deepcopy(train_accs)
+    self.train_acc5es = None
+    self.train_losses = copy.deepcopy(train_losses)
+    self.train_times  = None
+    self.arch_config  = copy.deepcopy(arch_config)
+    self.params     = params
+    self.flop       = flop
+    self.seed       = seed
+    self.epochs     = epochs
+    self.latency    = latency
+    # evaluation results
+    self.reset_eval()
+
+  def update_train_info(self, train_acc1es, train_acc5es, train_losses, train_times) -> None:
+    self.train_acc1es = train_acc1es
+    self.train_acc5es = train_acc5es
+    self.train_losses = train_losses
+    self.train_times  = train_times
+
+  def reset_pseudo_train_times(self, estimated_per_epoch_time: float) -> None:
+    """Assign the training times."""
+    train_times = OrderedDict()
+    for i in range(self.epochs):
+      train_times[i] = estimated_per_epoch_time
+    self.train_times = train_times
+
+  def reset_pseudo_eval_times(self, eval_name: Text, estimated_per_epoch_time: float) -> None:
+    """Assign the evaluation times."""
+    if eval_name not in self.eval_names: raise ValueError('invalid eval name : {:}'.format(eval_name))
+    for i in range(self.epochs):
+      self.eval_times['{:}@{:}'.format(eval_name,i)] = estimated_per_epoch_time
+
+  def reset_eval(self):
+    self.eval_names  = []
+    self.eval_acc1es = {}
+    self.eval_times  = {}
+    self.eval_losses = {}
+
+  def update_latency(self, latency):
+    self.latency = copy.deepcopy( latency )
+
+  def get_latency(self) -> float:
+    """Return the latency value in seconds. -1 represents not avaliable ; otherwise it should be a float value"""
+    if self.latency is None: return -1.0
+    else: return sum(self.latency) / len(self.latency)
+
+  def update_eval(self, accs, losses, times):  # new version
+    data_names = set([x.split('@')[0] for x in accs.keys()])
+    for data_name in data_names:
+      assert data_name not in self.eval_names, '{:} has already been added into eval-names'.format(data_name)
+      self.eval_names.append( data_name )
+      for iepoch in range(self.epochs):
+        xkey = '{:}@{:}'.format(data_name, iepoch)
+        self.eval_acc1es[ xkey ] = accs[ xkey ]
+        self.eval_losses[ xkey ] = losses[ xkey ]
+        self.eval_times [ xkey ] = times[ xkey ]
+
+  def update_OLD_eval(self, name, accs, losses): # old version
+    assert name not in self.eval_names, '{:} has already added'.format(name)
+    self.eval_names.append( name )
+    for iepoch in range(self.epochs):
+      if iepoch in accs:
+        self.eval_acc1es['{:}@{:}'.format(name,iepoch)] = accs[iepoch]
+        self.eval_losses['{:}@{:}'.format(name,iepoch)] = losses[iepoch]
+
+  def __repr__(self):
+    num_eval = len(self.eval_names)
+    set_name = '[' + ', '.join(self.eval_names) + ']'
+    return ('{name}({xname}, arch={arch}, FLOP={flop:.2f}M, Param={param:.3f}MB, seed={seed}, {num_eval} eval-sets: {set_name})'.format(name=self.__class__.__name__, xname=self.name, arch=self.arch_config['arch_str'], flop=self.flop, param=self.params, seed=self.seed, num_eval=num_eval, set_name=set_name))
+
+  def get_total_epoch(self):
+    return copy.deepcopy(self.epochs)
+
+  def get_times(self):
+    """Obtain the information regarding both training and evaluation time."""
+    if self.train_times is not None and isinstance(self.train_times, dict):
+      train_times = list( self.train_times.values() )
+      time_info = {'T-train@epoch': np.mean(train_times), 'T-train@total': np.sum(train_times)}
+    else:
+      time_info = {'T-train@epoch':                 None, 'T-train@total':               None }
+    for name in self.eval_names:
+      try:
+        xtimes = [self.eval_times['{:}@{:}'.format(name,i)] for i in range(self.epochs)]
+        time_info['T-{:}@epoch'.format(name)] = np.mean(xtimes)
+        time_info['T-{:}@total'.format(name)] = np.sum(xtimes)
+      except:
+        time_info['T-{:}@epoch'.format(name)] = None
+        time_info['T-{:}@total'.format(name)] = None
+    return time_info
+
+  def get_eval_set(self):
+    return self.eval_names
+
+  # get the training information
+  def get_train(self, iepoch=None):
+    if iepoch is None: iepoch = self.epochs-1
+    assert 0 <= iepoch < self.epochs, 'invalid iepoch={:} < {:}'.format(iepoch, self.epochs)
+    if self.train_times is not None:
+      xtime = self.train_times[iepoch]
+      atime = sum([self.train_times[i] for i in range(iepoch+1)])
+    else: xtime, atime = None, None
+    return {'iepoch'  : iepoch,
+            'loss'    : self.train_losses[iepoch],
+            'accuracy': self.train_acc1es[iepoch],
+            'cur_time': xtime,
+            'all_time': atime}
+
+  def get_eval(self, name, iepoch=None):
+    """Get the evaluation information ; there could be multiple evaluation sets (identified by the 'name' argument)."""
+    if iepoch is None: iepoch = self.epochs-1
+    assert 0 <= iepoch < self.epochs, 'invalid iepoch={:} < {:}'.format(iepoch, self.epochs)
+    def _internal_query(xname):
+      if isinstance(self.eval_times,dict) and len(self.eval_times) > 0:
+        xtime = self.eval_times['{:}@{:}'.format(xname, iepoch)]
+        atime = sum([self.eval_times['{:}@{:}'.format(xname, i)] for i in range(iepoch+1)])
+      else:
+        xtime, atime = None, None
+      return {'iepoch'  : iepoch,
+              'loss'    : self.eval_losses['{:}@{:}'.format(xname, iepoch)],
+              'accuracy': self.eval_acc1es['{:}@{:}'.format(xname, iepoch)],
+              'cur_time': xtime,
+              'all_time': atime}
+    if name == 'valid':
+      return _internal_query('x-valid')
+    else:
+      return _internal_query(name)
+
+  def get_net_param(self, clone=False):
+    if clone: return copy.deepcopy(self.net_state_dict)
+    else: return self.net_state_dict
+
+  def get_config(self, str2structure):
+    """This function is used to obtain the config dict for this architecture."""
+    if str2structure is None:
+      # In this case, this is architecture in the size search space of NATS-BENCH.
+      if 'name' in self.arch_config and self.arch_config['name'] == 'infer.shape.tiny':
+        return {'name': 'infer.shape.tiny', 'channels': self.arch_config['channels'],
+                'genotype': self.arch_config['genotype'], 'num_classes': self.arch_config['class_num']}
+      # In this case, this is architecture in the topology search space of NATS-BENCH.
+      else:
+        return {'name': 'infer.tiny', 'C': self.arch_config['channel'],
+                'N'   : self.arch_config['num_cells'],
+                'arch_str': self.arch_config['arch_str'], 'num_classes': self.arch_config['class_num']}
+    else:
+      # In this case, this is architecture in the size search space of NATS-BENCH.
+      if 'name' in self.arch_config and self.arch_config['name'] == 'infer.shape.tiny':
+        return {'name': 'infer.shape.tiny', 'channels': self.arch_config['channels'],
+                'genotype': str2structure(self.arch_config['genotype']), 'num_classes': self.arch_config['class_num']}
+      # In this case, this is architecture in the topology search space of NATS-BENCH.
+      else:
+        return {'name': 'infer.tiny', 'C': self.arch_config['channel'],
+                'N'   : self.arch_config['num_cells'],
+                'genotype': str2structure(self.arch_config['arch_str']), 'num_classes': self.arch_config['class_num']}
+
+  def state_dict(self):
+    _state_dict = {key: value for key, value in self.__dict__.items()}
+    return _state_dict
+
+  def load_state_dict(self, state_dict):
+    self.__dict__.update(state_dict)
+
+  @staticmethod
+  def create_from_state_dict(state_dict):
+    x = ResultsCount(None, None, None, None, None, None, None, None, None, None)
+    x.load_state_dict(state_dict)
+    return x