Add int search space

exps/NATS-algos/bohb.py

@@ -36,7 +36,9 @@ def get_topology_config_space(search_space, max_nodes=4):
     for i in range(1, max_nodes):
         for j in range(i):
             node_str = "{:}<-{:}".format(i, j)
-            cs.add_hyperparameter(ConfigSpace.CategoricalHyperparameter(node_str, search_space))
+            cs.add_hyperparameter(
+                ConfigSpace.CategoricalHyperparameter(node_str, search_space)
+            )
     return cs
 
 
@@ -44,7 +46,9 @@ def get_size_config_space(search_space):
     cs = ConfigSpace.ConfigurationSpace()
     for ilayer in range(search_space["numbers"]):
         node_str = "layer-{:}".format(ilayer)
-        cs.add_hyperparameter(ConfigSpace.CategoricalHyperparameter(node_str, search_space["candidates"]))
+        cs.add_hyperparameter(
+            ConfigSpace.CategoricalHyperparameter(node_str, search_space["candidates"])
+        )
     return cs
 
 
@@ -159,8 +163,14 @@ def main(xargs, api):
         current_best_index.append(api.query_index_by_arch(arch))
 
     best_arch = max(workers[0].trajectory, key=lambda x: x[0])[1]
-    logger.log("Best found configuration: {:} within {:.3f} s".format(best_arch, workers[0].total_times[-1]))
-    info = api.query_info_str_by_arch(best_arch, "200" if xargs.search_space == "tss" else "90")
+    logger.log(
+        "Best found configuration: {:} within {:.3f} s".format(
+            best_arch, workers[0].total_times[-1]
+        )
+    )
+    info = api.query_info_str_by_arch(
+        best_arch, "200" if xargs.search_space == "tss" else "90"
+    )
     logger.log("{:}".format(info))
     logger.log("-" * 100)
     logger.close()
@@ -169,7 +179,9 @@ def main(xargs, api):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser("BOHB: Robust and Efficient Hyperparameter Optimization at Scale")
+    parser = argparse.ArgumentParser(
+        "BOHB: Robust and Efficient Hyperparameter Optimization at Scale"
+    )
     parser.add_argument(
         "--dataset",
         type=str,
@@ -177,35 +189,80 @@ if __name__ == "__main__":
         help="Choose between Cifar10/100 and ImageNet-16.",
     )
     # general arg
-    parser.add_argument("--search_space", type=str, choices=["tss", "sss"], help="Choose the search space.")
     parser.add_argument(
-        "--time_budget", type=int, default=20000, help="The total time cost budge for searching (in seconds)."
+        "--search_space",
+        type=str,
+        choices=["tss", "sss"],
+        help="Choose the search space.",
+    )
+    parser.add_argument(
+        "--time_budget",
+        type=int,
+        default=20000,
+        help="The total time cost budge for searching (in seconds).",
+    )
+    parser.add_argument(
+        "--loops_if_rand", type=int, default=500, help="The total runs for evaluation."
     )
-    parser.add_argument("--loops_if_rand", type=int, default=500, help="The total runs for evaluation.")
     # BOHB
     parser.add_argument(
-        "--strategy", default="sampling", type=str, nargs="?", help="optimization strategy for the acquisition function"
-    )
-    parser.add_argument("--min_bandwidth", default=0.3, type=float, nargs="?", help="minimum bandwidth for KDE")
-    parser.add_argument(
-        "--num_samples", default=64, type=int, nargs="?", help="number of samples for the acquisition function"
+        "--strategy",
+        default="sampling",
+        type=str,
+        nargs="?",
+        help="optimization strategy for the acquisition function",
     )
     parser.add_argument(
-        "--random_fraction", default=0.33, type=float, nargs="?", help="fraction of random configurations"
+        "--min_bandwidth",
+        default=0.3,
+        type=float,
+        nargs="?",
+        help="minimum bandwidth for KDE",
     )
-    parser.add_argument("--bandwidth_factor", default=3, type=int, nargs="?", help="factor multiplied to the bandwidth")
     parser.add_argument(
-        "--n_iters", default=300, type=int, nargs="?", help="number of iterations for optimization method"
+        "--num_samples",
+        default=64,
+        type=int,
+        nargs="?",
+        help="number of samples for the acquisition function",
+    )
+    parser.add_argument(
+        "--random_fraction",
+        default=0.33,
+        type=float,
+        nargs="?",
+        help="fraction of random configurations",
+    )
+    parser.add_argument(
+        "--bandwidth_factor",
+        default=3,
+        type=int,
+        nargs="?",
+        help="factor multiplied to the bandwidth",
+    )
+    parser.add_argument(
+        "--n_iters",
+        default=300,
+        type=int,
+        nargs="?",
+        help="number of iterations for optimization method",
     )
     # log
-    parser.add_argument("--save_dir", type=str, default="./output/search", help="Folder to save checkpoints and log.")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="./output/search",
+        help="Folder to save checkpoints and log.",
+    )
     parser.add_argument("--rand_seed", type=int, default=-1, help="manual seed")
     args = parser.parse_args()
 
     api = create(None, args.search_space, fast_mode=False, verbose=False)
 
     args.save_dir = os.path.join(
-        "{:}-{:}".format(args.save_dir, args.search_space), "{:}-T{:}".format(args.dataset, args.time_budget), "BOHB"
+        "{:}-{:}".format(args.save_dir, args.search_space),
+        "{:}-T{:}".format(args.dataset, args.time_budget),
+        "BOHB",
     )
     print("save-dir : {:}".format(args.save_dir))
 

exps/NATS-algos/random_wo_share.py

@@ -19,7 +19,13 @@ if str(lib_dir) not in sys.path:
     sys.path.insert(0, str(lib_dir))
 from config_utils import load_config, dict2config, configure2str
 from datasets import get_datasets, SearchDataset
-from procedures import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint, get_optim_scheduler
+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 models import get_search_spaces
@@ -45,12 +51,16 @@ def main(xargs, api):
     current_best_index = []
     while len(total_time_cost) == 0 or total_time_cost[-1] < xargs.time_budget:
         arch = random_arch()
-        accuracy, _, _, total_cost = api.simulate_train_eval(arch, xargs.dataset, hp="12")
+        accuracy, _, _, total_cost = api.simulate_train_eval(
+            arch, xargs.dataset, hp="12"
+        )
         total_time_cost.append(total_cost)
         history.append(arch)
         if best_arch is None or best_acc < accuracy:
             best_acc, best_arch = accuracy, arch
-        logger.log("[{:03d}] : {:} : accuracy = {:.2f}%".format(len(history), arch, accuracy))
+        logger.log(
+            "[{:03d}] : {:} : accuracy = {:.2f}%".format(len(history), arch, accuracy)
+        )
         current_best_index.append(api.query_index_by_arch(best_arch))
     logger.log(
         "{:} best arch is {:}, accuracy = {:.2f}%, visit {:} archs with {:.1f} s.".format(
@@ -58,7 +68,9 @@ def main(xargs, api):
         )
     )
 
-    info = api.query_info_str_by_arch(best_arch, "200" if xargs.search_space == "tss" else "90")
+    info = api.query_info_str_by_arch(
+        best_arch, "200" if xargs.search_space == "tss" else "90"
+    )
     logger.log("{:}".format(info))
     logger.log("-" * 100)
     logger.close()
@@ -73,21 +85,38 @@ if __name__ == "__main__":
         choices=["cifar10", "cifar100", "ImageNet16-120"],
         help="Choose between Cifar10/100 and ImageNet-16.",
     )
-    parser.add_argument("--search_space", type=str, choices=["tss", "sss"], help="Choose the search space.")
+    parser.add_argument(
+        "--search_space",
+        type=str,
+        choices=["tss", "sss"],
+        help="Choose the search space.",
+    )
 
     parser.add_argument(
-        "--time_budget", type=int, default=20000, help="The total time cost budge for searching (in seconds)."
+        "--time_budget",
+        type=int,
+        default=20000,
+        help="The total time cost budge for searching (in seconds).",
+    )
+    parser.add_argument(
+        "--loops_if_rand", type=int, default=500, help="The total runs for evaluation."
     )
-    parser.add_argument("--loops_if_rand", type=int, default=500, help="The total runs for evaluation.")
     # log
-    parser.add_argument("--save_dir", type=str, default="./output/search", help="Folder to save checkpoints and log.")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="./output/search",
+        help="Folder to save checkpoints and log.",
+    )
     parser.add_argument("--rand_seed", type=int, default=-1, help="manual seed")
     args = parser.parse_args()
 
     api = create(None, args.search_space, fast_mode=True, verbose=False)
 
     args.save_dir = os.path.join(
-        "{:}-{:}".format(args.save_dir, args.search_space), "{:}-T{:}".format(args.dataset, args.time_budget), "RANDOM"
+        "{:}-{:}".format(args.save_dir, args.search_space),
+        "{:}-T{:}".format(args.dataset, args.time_budget),
+        "RANDOM",
     )
     print("save-dir : {:}".format(args.save_dir))
 

exps/NATS-algos/regularized_ea.py

@@ -23,7 +23,13 @@ if str(lib_dir) not in sys.path:
     sys.path.insert(0, str(lib_dir))
 from config_utils import load_config, dict2config, configure2str
 from datasets import get_datasets, SearchDataset
-from procedures import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint, get_optim_scheduler
+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 models import CellStructure, get_search_spaces
@@ -103,7 +109,15 @@ def mutate_size_func(info):
 
 
 def regularized_evolution(
-    cycles, population_size, sample_size, time_budget, random_arch, mutate_arch, api, use_proxy, dataset
+    cycles,
+    population_size,
+    sample_size,
+    time_budget,
+    random_arch,
+    mutate_arch,
+    api,
+    use_proxy,
+    dataset,
 ):
     """Algorithm for regularized evolution (i.e. aging evolution).
 
@@ -122,7 +136,10 @@ def regularized_evolution(
     """
     population = collections.deque()
     api.reset_time()
-    history, total_time_cost = [], []  # Not used by the algorithm, only used to report results.
+    history, total_time_cost = (
+        [],
+        [],
+    )  # Not used by the algorithm, only used to report results.
     current_best_index = []
     # Initialize the population with random models.
     while len(population) < population_size:
@@ -135,7 +152,9 @@ def regularized_evolution(
         population.append(model)
         history.append((model.accuracy, model.arch))
         total_time_cost.append(total_cost)
-        current_best_index.append(api.query_index_by_arch(max(history, key=lambda x: x[0])[1]))
+        current_best_index.append(
+            api.query_index_by_arch(max(history, key=lambda x: x[0])[1])
+        )
 
     # Carry out evolution in cycles. Each cycle produces a model and removes another.
     while total_time_cost[-1] < time_budget:
@@ -160,7 +179,9 @@ def regularized_evolution(
         # Append the info
         population.append(child)
         history.append((child.accuracy, child.arch))
-        current_best_index.append(api.query_index_by_arch(max(history, key=lambda x: x[0])[1]))
+        current_best_index.append(
+            api.query_index_by_arch(max(history, key=lambda x: x[0])[1])
+        )
         total_time_cost.append(total_cost)
 
         # Remove the oldest model.
@@ -183,7 +204,10 @@ def main(xargs, api):
 
     x_start_time = time.time()
     logger.log("{:} use api : {:}".format(time_string(), api))
-    logger.log("-" * 30 + " start searching with the time budget of {:} s".format(xargs.time_budget))
+    logger.log(
+        "-" * 30
+        + " start searching with the time budget of {:} s".format(xargs.time_budget)
+    )
     history, current_best_index, total_times = regularized_evolution(
         xargs.ea_cycles,
         xargs.ea_population,
@@ -203,7 +227,9 @@ def main(xargs, api):
     best_arch = max(history, key=lambda x: x[0])[1]
     logger.log("{:} best arch is {:}".format(time_string(), best_arch))
 
-    info = api.query_info_str_by_arch(best_arch, "200" if xargs.search_space == "tss" else "90")
+    info = api.query_info_str_by_arch(
+        best_arch, "200" if xargs.search_space == "tss" else "90"
+    )
     logger.log("{:}".format(info))
     logger.log("-" * 100)
     logger.close()
@@ -218,19 +244,39 @@ if __name__ == "__main__":
         choices=["cifar10", "cifar100", "ImageNet16-120"],
         help="Choose between Cifar10/100 and ImageNet-16.",
     )
-    parser.add_argument("--search_space", type=str, choices=["tss", "sss"], help="Choose the search space.")
+    parser.add_argument(
+        "--search_space",
+        type=str,
+        choices=["tss", "sss"],
+        help="Choose the search space.",
+    )
     # hyperparameters for REA
     parser.add_argument("--ea_cycles", type=int, help="The number of cycles in EA.")
     parser.add_argument("--ea_population", type=int, help="The population size in EA.")
     parser.add_argument("--ea_sample_size", type=int, help="The sample size in EA.")
     parser.add_argument(
-        "--time_budget", type=int, default=20000, help="The total time cost budge for searching (in seconds)."
+        "--time_budget",
+        type=int,
+        default=20000,
+        help="The total time cost budge for searching (in seconds).",
+    )
+    parser.add_argument(
+        "--use_proxy",
+        type=int,
+        default=1,
+        help="Whether to use the proxy (H0) task or not.",
     )
-    parser.add_argument("--use_proxy", type=int, default=1, help="Whether to use the proxy (H0) task or not.")
     #
-    parser.add_argument("--loops_if_rand", type=int, default=500, help="The total runs for evaluation.")
+    parser.add_argument(
+        "--loops_if_rand", type=int, default=500, help="The total runs for evaluation."
+    )
     # log
-    parser.add_argument("--save_dir", type=str, default="./output/search", help="Folder to save checkpoints and log.")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="./output/search",
+        help="Folder to save checkpoints and log.",
+    )
     parser.add_argument("--rand_seed", type=int, default=-1, help="manual seed")
     args = parser.parse_args()
 
@@ -238,7 +284,9 @@ if __name__ == "__main__":
 
     args.save_dir = os.path.join(
         "{:}-{:}".format(args.save_dir, args.search_space),
-        "{:}-T{:}{:}".format(args.dataset, args.time_budget, "" if args.use_proxy > 0 else "-FULL"),
+        "{:}-T{:}{:}".format(
+            args.dataset, args.time_budget, "" if args.use_proxy > 0 else "-FULL"
+        ),
         "R-EA-SS{:}".format(args.ea_sample_size),
     )
     print("save-dir : {:}".format(args.save_dir))

exps/NATS-algos/reinforce.py

@@ -23,7 +23,13 @@ if str(lib_dir) not in sys.path:
     sys.path.insert(0, str(lib_dir))
 from config_utils import load_config, dict2config, configure2str
 from datasets import get_datasets, SearchDataset
-from procedures import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint, get_optim_scheduler
+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 models import CellStructure, get_search_spaces
@@ -40,7 +46,9 @@ class PolicyTopology(nn.Module):
             for j in range(i):
                 node_str = "{:}<-{:}".format(i, j)
                 self.edge2index[node_str] = len(self.edge2index)
-        self.arch_parameters = nn.Parameter(1e-3 * torch.randn(len(self.edge2index), len(search_space)))
+        self.arch_parameters = nn.Parameter(
+            1e-3 * torch.randn(len(self.edge2index), len(search_space))
+        )
 
     def generate_arch(self, actions):
         genotypes = []
@@ -76,7 +84,9 @@ class PolicySize(nn.Module):
         super(PolicySize, self).__init__()
         self.candidates = search_space["candidates"]
         self.numbers = search_space["numbers"]
-        self.arch_parameters = nn.Parameter(1e-3 * torch.randn(self.numbers, len(self.candidates)))
+        self.arch_parameters = nn.Parameter(
+            1e-3 * torch.randn(self.numbers, len(self.candidates))
+        )
 
     def generate_arch(self, actions):
         channels = [str(self.candidates[i]) for i in actions]
@@ -103,7 +113,9 @@ class ExponentialMovingAverage(object):
         self._momentum = momentum
 
     def update(self, value):
-        self._numerator = self._momentum * self._numerator + (1 - self._momentum) * value
+        self._numerator = (
+            self._momentum * self._numerator + (1 - self._momentum) * value
+        )
         self._denominator = self._momentum * self._denominator + (1 - self._momentum)
 
     def value(self):
@@ -143,14 +155,18 @@ def main(xargs, api):
 
     # REINFORCE
     x_start_time = time.time()
-    logger.log("Will start searching with time budget of {:} s.".format(xargs.time_budget))
+    logger.log(
+        "Will start searching with time budget of {:} s.".format(xargs.time_budget)
+    )
     total_steps, total_costs, trace = 0, [], []
     current_best_index = []
     while len(total_costs) == 0 or total_costs[-1] < xargs.time_budget:
         start_time = time.time()
         log_prob, action = select_action(policy)
         arch = policy.generate_arch(action)
-        reward, _, _, current_total_cost = api.simulate_train_eval(arch, xargs.dataset, hp="12")
+        reward, _, _, current_total_cost = api.simulate_train_eval(
+            arch, xargs.dataset, hp="12"
+        )
         trace.append((reward, arch))
         total_costs.append(current_total_cost)
 
@@ -168,7 +184,9 @@ def main(xargs, api):
             )
         )
         # to analyze
-        current_best_index.append(api.query_index_by_arch(max(trace, key=lambda x: x[0])[1]))
+        current_best_index.append(
+            api.query_index_by_arch(max(trace, key=lambda x: x[0])[1])
+        )
     # best_arch = policy.genotype() # first version
     best_arch = max(trace, key=lambda x: x[0])[1]
     logger.log(
@@ -176,7 +194,9 @@ def main(xargs, api):
             total_steps, total_costs[-1], time.time() - x_start_time
         )
     )
-    info = api.query_info_str_by_arch(best_arch, "200" if xargs.search_space == "tss" else "90")
+    info = api.query_info_str_by_arch(
+        best_arch, "200" if xargs.search_space == "tss" else "90"
+    )
     logger.log("{:}".format(info))
     logger.log("-" * 100)
     logger.close()
@@ -193,17 +213,38 @@ if __name__ == "__main__":
         choices=["cifar10", "cifar100", "ImageNet16-120"],
         help="Choose between Cifar10/100 and ImageNet-16.",
     )
-    parser.add_argument("--search_space", type=str, choices=["tss", "sss"], help="Choose the search space.")
-    parser.add_argument("--learning_rate", type=float, help="The learning rate for REINFORCE.")
-    parser.add_argument("--EMA_momentum", type=float, default=0.9, help="The momentum value for EMA.")
     parser.add_argument(
-        "--time_budget", type=int, default=20000, help="The total time cost budge for searching (in seconds)."
+        "--search_space",
+        type=str,
+        choices=["tss", "sss"],
+        help="Choose the search space.",
     )
-    parser.add_argument("--loops_if_rand", type=int, default=500, help="The total runs for evaluation.")
-    # log
-    parser.add_argument("--save_dir", type=str, default="./output/search", help="Folder to save checkpoints and log.")
     parser.add_argument(
-        "--arch_nas_dataset", type=str, help="The path to load the architecture dataset (tiny-nas-benchmark)."
+        "--learning_rate", type=float, help="The learning rate for REINFORCE."
+    )
+    parser.add_argument(
+        "--EMA_momentum", type=float, default=0.9, help="The momentum value for EMA."
+    )
+    parser.add_argument(
+        "--time_budget",
+        type=int,
+        default=20000,
+        help="The total time cost budge for searching (in seconds).",
+    )
+    parser.add_argument(
+        "--loops_if_rand", type=int, default=500, help="The total runs for evaluation."
+    )
+    # log
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="./output/search",
+        help="Folder to save checkpoints and log.",
+    )
+    parser.add_argument(
+        "--arch_nas_dataset",
+        type=str,
+        help="The path to load the architecture dataset (tiny-nas-benchmark).",
     )
     parser.add_argument("--print_freq", type=int, help="print frequency (default: 200)")
     parser.add_argument("--rand_seed", type=int, default=-1, help="manual seed")

exps/NATS-algos/search-cell.py

@@ -37,7 +37,13 @@ if str(lib_dir) not in sys.path:
     sys.path.insert(0, str(lib_dir))
 from config_utils import load_config, dict2config, configure2str
 from datasets import get_datasets, get_nas_search_loaders
-from procedures import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint, get_optim_scheduler
+from procedures import (
+    prepare_seed,
+    prepare_logger,
+    save_checkpoint,
+    copy_checkpoint,
+    get_optim_scheduler,
+)
 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
@@ -49,7 +55,9 @@ def _concat(xs):
     return torch.cat([x.view(-1) for x in xs])
 
 
-def _hessian_vector_product(vector, network, criterion, base_inputs, base_targets, r=1e-2):
+def _hessian_vector_product(
+    vector, network, criterion, base_inputs, base_targets, r=1e-2
+):
     R = r / _concat(vector).norm()
     for p, v in zip(network.weights, vector):
         p.data.add_(R, v)
@@ -68,7 +76,15 @@ def _hessian_vector_product(vector, network, criterion, base_inputs, base_target
     return [(x - y).div_(2 * R) for x, y in zip(grads_p, grads_n)]
 
 
-def backward_step_unrolled(network, criterion, base_inputs, base_targets, w_optimizer, arch_inputs, arch_targets):
+def backward_step_unrolled(
+    network,
+    criterion,
+    base_inputs,
+    base_targets,
+    w_optimizer,
+    arch_inputs,
+    arch_targets,
+):
     # _compute_unrolled_model
     _, logits = network(base_inputs)
     loss = criterion(logits, base_targets)
@@ -80,7 +96,9 @@ def backward_step_unrolled(network, criterion, base_inputs, base_targets, w_opti
     with torch.no_grad():
         theta = _concat(network.weights)
         try:
-            moment = _concat(w_optimizer.state[v]["momentum_buffer"] for v in network.weights)
+            moment = _concat(
+                w_optimizer.state[v]["momentum_buffer"] for v in network.weights
+            )
             moment = moment.mul_(momentum)
         except:
             moment = torch.zeros_like(theta)
@@ -105,7 +123,9 @@ def backward_step_unrolled(network, criterion, base_inputs, base_targets, w_opti
 
     dalpha = unrolled_model.arch_parameters.grad
     vector = [v.grad.data for v in unrolled_model.weights]
-    [implicit_grads] = _hessian_vector_product(vector, network, criterion, base_inputs, base_targets)
+    [implicit_grads] = _hessian_vector_product(
+        vector, network, criterion, base_inputs, base_targets
+    )
 
     dalpha.data.sub_(LR, implicit_grads.data)
 
@@ -116,13 +136,26 @@ def backward_step_unrolled(network, criterion, base_inputs, base_targets, w_opti
     return unrolled_loss.detach(), unrolled_logits.detach()
 
 
-def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer, epoch_str, print_freq, algo, logger):
+def search_func(
+    xloader,
+    network,
+    criterion,
+    scheduler,
+    w_optimizer,
+    a_optimizer,
+    epoch_str,
+    print_freq,
+    algo,
+    logger,
+):
     data_time, batch_time = AverageMeter(), AverageMeter()
     base_losses, base_top1, base_top5 = AverageMeter(), AverageMeter(), AverageMeter()
     arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter()
     end = time.time()
     network.train()
-    for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(xloader):
+    for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(
+        xloader
+    ):
         scheduler.update(None, 1.0 * step / len(xloader))
         base_inputs = base_inputs.cuda(non_blocking=True)
         arch_inputs = arch_inputs.cuda(non_blocking=True)
@@ -155,7 +188,9 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer
         base_loss.backward()
         w_optimizer.step()
         # record
-        base_prec1, base_prec5 = obtain_accuracy(logits.data, base_targets.data, topk=(1, 5))
+        base_prec1, base_prec5 = obtain_accuracy(
+            logits.data, base_targets.data, topk=(1, 5)
+        )
         base_losses.update(base_loss.item(), base_inputs.size(0))
         base_top1.update(base_prec1.item(), base_inputs.size(0))
         base_top5.update(base_prec5.item(), base_inputs.size(0))
@@ -174,7 +209,13 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer
         network.zero_grad()
         if algo == "darts-v2":
             arch_loss, logits = backward_step_unrolled(
-                network, criterion, base_inputs, base_targets, w_optimizer, arch_inputs, arch_targets
+                network,
+                criterion,
+                base_inputs,
+                base_targets,
+                w_optimizer,
+                arch_inputs,
+                arch_targets,
             )
             a_optimizer.step()
         elif algo == "random" or algo == "enas":
@@ -187,7 +228,9 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer
             arch_loss.backward()
             a_optimizer.step()
         # record
-        arch_prec1, arch_prec5 = obtain_accuracy(logits.data, arch_targets.data, topk=(1, 5))
+        arch_prec1, arch_prec5 = obtain_accuracy(
+            logits.data, arch_targets.data, topk=(1, 5)
+        )
         arch_losses.update(arch_loss.item(), arch_inputs.size(0))
         arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
         arch_top5.update(arch_prec5.item(), arch_inputs.size(0))
@@ -197,7 +240,11 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer
         end = time.time()
 
         if step % print_freq == 0 or step + 1 == len(xloader):
-            Sstr = "*SEARCH* " + time_string() + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader))
+            Sstr = (
+                "*SEARCH* "
+                + time_string()
+                + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader))
+            )
             Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
                 batch_time=batch_time, data_time=data_time
             )
@@ -208,14 +255,31 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer
                 loss=arch_losses, top1=arch_top1, top5=arch_top5
             )
             logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Astr)
-    return base_losses.avg, base_top1.avg, base_top5.avg, arch_losses.avg, arch_top1.avg, arch_top5.avg
+    return (
+        base_losses.avg,
+        base_top1.avg,
+        base_top5.avg,
+        arch_losses.avg,
+        arch_top1.avg,
+        arch_top5.avg,
+    )
 
 
-def train_controller(xloader, network, criterion, optimizer, prev_baseline, epoch_str, print_freq, logger):
+def train_controller(
+    xloader, network, criterion, optimizer, prev_baseline, epoch_str, print_freq, logger
+):
     # config. (containing some necessary arg)
     #   baseline: The baseline score (i.e. average val_acc) from the previous epoch
     data_time, batch_time = AverageMeter(), AverageMeter()
-    GradnormMeter, LossMeter, ValAccMeter, EntropyMeter, BaselineMeter, RewardMeter, xend = (
+    (
+        GradnormMeter,
+        LossMeter,
+        ValAccMeter,
+        EntropyMeter,
+        BaselineMeter,
+        RewardMeter,
+        xend,
+    ) = (
         AverageMeter(),
         AverageMeter(),
         AverageMeter(),
@@ -255,7 +319,9 @@ def train_controller(xloader, network, criterion, optimizer, prev_baseline, epoc
         if prev_baseline is None:
             baseline = val_top1
         else:
-            baseline = prev_baseline - (1 - controller_bl_dec) * (prev_baseline - reward)
+            baseline = prev_baseline - (1 - controller_bl_dec) * (
+                prev_baseline - reward
+            )
 
         loss = -1 * log_prob * (reward - baseline)
 
@@ -274,7 +340,9 @@ def train_controller(xloader, network, criterion, optimizer, prev_baseline, epoc
         batch_time.update(time.time() - xend)
         xend = time.time()
         if (step + 1) % controller_num_aggregate == 0:
-            grad_norm = torch.nn.utils.clip_grad_norm_(network.controller.parameters(), 5.0)
+            grad_norm = torch.nn.utils.clip_grad_norm_(
+                network.controller.parameters(), 5.0
+            )
             GradnormMeter.update(grad_norm)
             optimizer.step()
             network.controller.zero_grad()
@@ -283,13 +351,18 @@ def train_controller(xloader, network, criterion, optimizer, prev_baseline, epoc
             Sstr = (
                 "*Train-Controller* "
                 + time_string()
-                + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, controller_train_steps * controller_num_aggregate)
+                + " [{:}][{:03d}/{:03d}]".format(
+                    epoch_str, step, controller_train_steps * controller_num_aggregate
+                )
             )
             Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
                 batch_time=batch_time, data_time=data_time
             )
             Wstr = "[Loss {loss.val:.3f} ({loss.avg:.3f})  Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Reward {reward.val:.2f} ({reward.avg:.2f})] Baseline {basel.val:.2f} ({basel.avg:.2f})".format(
-                loss=LossMeter, top1=ValAccMeter, reward=RewardMeter, basel=BaselineMeter
+                loss=LossMeter,
+                top1=ValAccMeter,
+                reward=RewardMeter,
+                basel=BaselineMeter,
             )
             Estr = "Entropy={:.4f} ({:.4f})".format(EntropyMeter.val, EntropyMeter.avg)
             logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Estr)
@@ -323,7 +396,9 @@ def get_best_arch(xloader, network, n_samples, algo):
                 loader_iter = iter(xloader)
                 inputs, targets = next(loader_iter)
             _, logits = network(inputs.cuda(non_blocking=True))
-            val_top1, val_top5 = obtain_accuracy(logits.cpu().data, targets.data, topk=(1, 5))
+            val_top1, val_top5 = obtain_accuracy(
+                logits.cpu().data, targets.data, topk=(1, 5)
+            )
             valid_accs.append(val_top1.item())
         best_idx = np.argmax(valid_accs)
         best_arch, best_valid_acc = archs[best_idx], valid_accs[best_idx]
@@ -344,7 +419,9 @@ def valid_func(xloader, network, criterion, algo, logger):
             _, logits = network(arch_inputs.cuda(non_blocking=True))
             arch_loss = criterion(logits, arch_targets)
             # record
-            arch_prec1, arch_prec5 = obtain_accuracy(logits.data, arch_targets.data, topk=(1, 5))
+            arch_prec1, arch_prec5 = obtain_accuracy(
+                logits.data, arch_targets.data, topk=(1, 5)
+            )
             arch_losses.update(arch_loss.item(), arch_inputs.size(0))
             arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
             arch_top5.update(arch_prec5.item(), arch_inputs.size(0))
@@ -363,11 +440,17 @@ def main(xargs):
     prepare_seed(xargs.rand_seed)
     logger = prepare_logger(args)
 
-    train_data, valid_data, xshape, class_num = get_datasets(xargs.dataset, xargs.data_path, -1)
+    train_data, valid_data, xshape, class_num = get_datasets(
+        xargs.dataset, xargs.data_path, -1
+    )
     if xargs.overwite_epochs is None:
         extra_info = {"class_num": class_num, "xshape": xshape}
     else:
-        extra_info = {"class_num": class_num, "xshape": xshape, "epochs": xargs.overwite_epochs}
+        extra_info = {
+            "class_num": class_num,
+            "xshape": xshape,
+            "epochs": xargs.overwite_epochs,
+        }
     config = load_config(xargs.config_path, extra_info, logger)
     search_loader, train_loader, valid_loader = get_nas_search_loaders(
         train_data,
@@ -405,7 +488,9 @@ def main(xargs):
     search_model.set_algo(xargs.algo)
     logger.log("{:}".format(search_model))
 
-    w_optimizer, w_scheduler, criterion = get_optim_scheduler(search_model.weights, config)
+    w_optimizer, w_scheduler, criterion = get_optim_scheduler(
+        search_model.weights, config
+    )
     a_optimizer = torch.optim.Adam(
         search_model.alphas,
         lr=xargs.arch_learning_rate,
@@ -426,13 +511,23 @@ def main(xargs):
         api = None
     logger.log("{:} create API = {:} done".format(time_string(), api))
 
-    last_info, model_base_path, model_best_path = logger.path("info"), logger.path("model"), logger.path("best")
+    last_info, model_base_path, model_best_path = (
+        logger.path("info"),
+        logger.path("model"),
+        logger.path("best"),
+    )
     network, criterion = search_model.cuda(), criterion.cuda()  # use a single GPU
 
-    last_info, model_base_path, model_best_path = logger.path("info"), logger.path("model"), logger.path("best")
+    last_info, model_base_path, model_best_path = (
+        logger.path("info"),
+        logger.path("model"),
+        logger.path("best"),
+    )
 
     if last_info.exists():  # automatically resume from previous checkpoint
-        logger.log("=> loading checkpoint of the last-info '{:}' start".format(last_info))
+        logger.log(
+            "=> loading checkpoint of the last-info '{:}' start".format(last_info)
+        )
         last_info = torch.load(last_info)
         start_epoch = last_info["epoch"]
         checkpoint = torch.load(last_info["last_checkpoint"])
@@ -444,11 +539,17 @@ def main(xargs):
         w_optimizer.load_state_dict(checkpoint["w_optimizer"])
         a_optimizer.load_state_dict(checkpoint["a_optimizer"])
         logger.log(
-            "=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(last_info, start_epoch)
+            "=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(
+                last_info, start_epoch
+            )
         )
     else:
         logger.log("=> do not find the last-info file : {:}".format(last_info))
-        start_epoch, valid_accuracies, genotypes = 0, {"best": -1}, {-1: network.return_topK(1, True)[0]}
+        start_epoch, valid_accuracies, genotypes = (
+            0,
+            {"best": -1},
+            {-1: network.return_topK(1, True)[0]},
+        )
         baseline = None
 
     # start training
@@ -460,15 +561,35 @@ def main(xargs):
     )
     for epoch in range(start_epoch, total_epoch):
         w_scheduler.update(epoch, 0.0)
-        need_time = "Time Left: {:}".format(convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
+        need_time = "Time Left: {:}".format(
+            convert_secs2time(epoch_time.val * (total_epoch - epoch), True)
+        )
         epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
-        logger.log("\n[Search the {:}-th epoch] {:}, LR={:}".format(epoch_str, need_time, min(w_scheduler.get_lr())))
+        logger.log(
+            "\n[Search the {:}-th epoch] {:}, LR={:}".format(
+                epoch_str, need_time, min(w_scheduler.get_lr())
+            )
+        )
 
         network.set_drop_path(float(epoch + 1) / total_epoch, xargs.drop_path_rate)
         if xargs.algo == "gdas":
-            network.set_tau(xargs.tau_max - (xargs.tau_max - xargs.tau_min) * epoch / (total_epoch - 1))
-            logger.log("[RESET tau as : {:} and drop_path as {:}]".format(network.tau, network.drop_path))
-        search_w_loss, search_w_top1, search_w_top5, search_a_loss, search_a_top1, search_a_top5 = search_func(
+            network.set_tau(
+                xargs.tau_max
+                - (xargs.tau_max - xargs.tau_min) * epoch / (total_epoch - 1)
+            )
+            logger.log(
+                "[RESET tau as : {:} and drop_path as {:}]".format(
+                    network.tau, network.drop_path
+                )
+            )
+        (
+            search_w_loss,
+            search_w_top1,
+            search_w_top5,
+            search_a_loss,
+            search_a_top1,
+            search_a_top5,
+        ) = search_func(
             search_loader,
             network,
             criterion,
@@ -493,7 +614,14 @@ def main(xargs):
         )
         if xargs.algo == "enas":
             ctl_loss, ctl_acc, baseline, ctl_reward = train_controller(
-                valid_loader, network, criterion, a_optimizer, baseline, epoch_str, xargs.print_freq, logger
+                valid_loader,
+                network,
+                criterion,
+                a_optimizer,
+                baseline,
+                epoch_str,
+                xargs.print_freq,
+                logger,
             )
             logger.log(
                 "[{:}] controller : loss={:}, acc={:}, baseline={:}, reward={:}".format(
@@ -501,7 +629,9 @@ def main(xargs):
                 )
             )
 
-        genotype, temp_accuracy = get_best_arch(valid_loader, network, xargs.eval_candidate_num, xargs.algo)
+        genotype, temp_accuracy = get_best_arch(
+            valid_loader, network, xargs.eval_candidate_num, xargs.algo
+        )
         if xargs.algo == "setn" or xargs.algo == "enas":
             network.set_cal_mode("dynamic", genotype)
         elif xargs.algo == "gdas":
@@ -512,8 +642,14 @@ def main(xargs):
             network.set_cal_mode("urs", None)
         else:
             raise ValueError("Invalid algorithm name : {:}".format(xargs.algo))
-        logger.log("[{:}] - [get_best_arch] : {:} -> {:}".format(epoch_str, genotype, temp_accuracy))
-        valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(valid_loader, network, criterion, xargs.algo, logger)
+        logger.log(
+            "[{:}] - [get_best_arch] : {:} -> {:}".format(
+                epoch_str, genotype, temp_accuracy
+            )
+        )
+        valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
+            valid_loader, network, criterion, xargs.algo, logger
+        )
         logger.log(
             "[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}% | {:}".format(
                 epoch_str, valid_a_loss, valid_a_top1, valid_a_top5, genotype
@@ -522,7 +658,9 @@ def main(xargs):
         valid_accuracies[epoch] = valid_a_top1
 
         genotypes[epoch] = genotype
-        logger.log("<<<--->>> The {:}-th epoch : {:}".format(epoch_str, genotypes[epoch]))
+        logger.log(
+            "<<<--->>> The {:}-th epoch : {:}".format(epoch_str, genotypes[epoch])
+        )
         # save checkpoint
         save_path = save_checkpoint(
             {
@@ -558,7 +696,9 @@ def main(xargs):
 
     # the final post procedure : count the time
     start_time = time.time()
-    genotype, temp_accuracy = get_best_arch(valid_loader, network, xargs.eval_candidate_num, xargs.algo)
+    genotype, temp_accuracy = get_best_arch(
+        valid_loader, network, xargs.eval_candidate_num, xargs.algo
+    )
     if xargs.algo == "setn" or xargs.algo == "enas":
         network.set_cal_mode("dynamic", genotype)
     elif xargs.algo == "gdas":
@@ -571,8 +711,14 @@ def main(xargs):
         raise ValueError("Invalid algorithm name : {:}".format(xargs.algo))
     search_time.update(time.time() - start_time)
 
-    valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(valid_loader, network, criterion, xargs.algo, logger)
-    logger.log("Last : the gentotype is : {:}, with the validation accuracy of {:.3f}%.".format(genotype, valid_a_top1))
+    valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
+        valid_loader, network, criterion, xargs.algo, logger
+    )
+    logger.log(
+        "Last : the gentotype is : {:}, with the validation accuracy of {:.3f}%.".format(
+            genotype, valid_a_top1
+        )
+    )
 
     logger.log("\n" + "-" * 100)
     # check the performance from the architecture dataset
@@ -595,7 +741,13 @@ if __name__ == "__main__":
         choices=["cifar10", "cifar100", "ImageNet16-120"],
         help="Choose between Cifar10/100 and ImageNet-16.",
     )
-    parser.add_argument("--search_space", type=str, default="tss", choices=["tss"], help="The search space name.")
+    parser.add_argument(
+        "--search_space",
+        type=str,
+        default="tss",
+        choices=["tss"],
+        help="The search space name.",
+    )
     parser.add_argument(
         "--algo",
         type=str,
@@ -603,18 +755,35 @@ if __name__ == "__main__":
         help="The search space name.",
     )
     parser.add_argument(
-        "--use_api", type=int, default=1, choices=[0, 1], help="Whether use API or not (which will cost much memory)."
+        "--use_api",
+        type=int,
+        default=1,
+        choices=[0, 1],
+        help="Whether use API or not (which will cost much memory).",
     )
     # FOR GDAS
-    parser.add_argument("--tau_min", type=float, default=0.1, help="The minimum tau for Gumbel Softmax.")
-    parser.add_argument("--tau_max", type=float, default=10, help="The maximum tau for Gumbel Softmax.")
+    parser.add_argument(
+        "--tau_min", type=float, default=0.1, help="The minimum tau for Gumbel Softmax."
+    )
+    parser.add_argument(
+        "--tau_max", type=float, default=10, help="The maximum tau for Gumbel Softmax."
+    )
     # channels and number-of-cells
-    parser.add_argument("--max_nodes", type=int, default=4, help="The maximum number of nodes.")
-    parser.add_argument("--channel", type=int, default=16, help="The number of channels.")
-    parser.add_argument("--num_cells", type=int, default=5, help="The number of cells in one stage.")
+    parser.add_argument(
+        "--max_nodes", type=int, default=4, help="The maximum number of nodes."
+    )
+    parser.add_argument(
+        "--channel", type=int, default=16, help="The number of channels."
+    )
+    parser.add_argument(
+        "--num_cells", type=int, default=5, help="The number of cells in one stage."
+    )
     #
     parser.add_argument(
-        "--eval_candidate_num", type=int, default=100, help="The number of selected architectures to evaluate."
+        "--eval_candidate_num",
+        type=int,
+        default=100,
+        help="The number of selected architectures to evaluate.",
     )
     #
     parser.add_argument(
@@ -625,7 +794,11 @@ if __name__ == "__main__":
         help="Whether use track_running_stats or not in the BN layer.",
     )
     parser.add_argument(
-        "--affine", type=int, default=0, choices=[0, 1], help="Whether use affine=True or False in the BN layer."
+        "--affine",
+        type=int,
+        default=0,
+        choices=[0, 1],
+        help="Whether use affine=True or False in the BN layer.",
     )
     parser.add_argument(
         "--config_path",
@@ -634,17 +807,43 @@ if __name__ == "__main__":
         help="The path of configuration.",
     )
     parser.add_argument(
-        "--overwite_epochs", type=int, help="The number of epochs to overwrite that value in config files."
+        "--overwite_epochs",
+        type=int,
+        help="The number of epochs to overwrite that value in config files.",
     )
     # architecture leraning rate
-    parser.add_argument("--arch_learning_rate", type=float, default=3e-4, help="learning rate for arch encoding")
-    parser.add_argument("--arch_weight_decay", type=float, default=1e-3, help="weight decay for arch encoding")
-    parser.add_argument("--arch_eps", type=float, default=1e-8, help="weight decay for arch encoding")
+    parser.add_argument(
+        "--arch_learning_rate",
+        type=float,
+        default=3e-4,
+        help="learning rate for arch encoding",
+    )
+    parser.add_argument(
+        "--arch_weight_decay",
+        type=float,
+        default=1e-3,
+        help="weight decay for arch encoding",
+    )
+    parser.add_argument(
+        "--arch_eps", type=float, default=1e-8, help="weight decay for arch encoding"
+    )
     parser.add_argument("--drop_path_rate", type=float, help="The drop path rate.")
     # log
-    parser.add_argument("--workers", type=int, default=2, help="number of data loading workers (default: 2)")
-    parser.add_argument("--save_dir", type=str, default="./output/search", help="Folder to save checkpoints and log.")
-    parser.add_argument("--print_freq", type=int, default=200, help="print frequency (default: 200)")
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=2,
+        help="number of data loading workers (default: 2)",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="./output/search",
+        help="Folder to save checkpoints and log.",
+    )
+    parser.add_argument(
+        "--print_freq", type=int, default=200, help="print frequency (default: 200)"
+    )
     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:
@@ -653,14 +852,20 @@ if __name__ == "__main__":
         args.save_dir = os.path.join(
             "{:}-{:}".format(args.save_dir, args.search_space),
             args.dataset,
-            "{:}-affine{:}_BN{:}-{:}".format(args.algo, args.affine, args.track_running_stats, args.drop_path_rate),
+            "{:}-affine{:}_BN{:}-{:}".format(
+                args.algo, args.affine, args.track_running_stats, args.drop_path_rate
+            ),
         )
     else:
         args.save_dir = os.path.join(
             "{:}-{:}".format(args.save_dir, args.search_space),
             args.dataset,
             "{:}-affine{:}_BN{:}-E{:}-{:}".format(
-                args.algo, args.affine, args.track_running_stats, args.overwite_epochs, args.drop_path_rate
+                args.algo,
+                args.affine,
+                args.track_running_stats,
+                args.overwite_epochs,
+                args.drop_path_rate,
             ),
         )
 

exps/NATS-algos/search-size.py

@@ -38,7 +38,13 @@ if str(lib_dir) not in sys.path:
     sys.path.insert(0, str(lib_dir))
 from config_utils import load_config, dict2config, configure2str
 from datasets import get_datasets, get_nas_search_loaders
-from procedures import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint, get_optim_scheduler
+from procedures import (
+    prepare_seed,
+    prepare_logger,
+    save_checkpoint,
+    copy_checkpoint,
+    get_optim_scheduler,
+)
 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
@@ -55,7 +61,9 @@ class ExponentialMovingAverage(object):
         self._momentum = momentum
 
     def update(self, value):
-        self._numerator = self._momentum * self._numerator + (1 - self._momentum) * value
+        self._numerator = (
+            self._momentum * self._numerator + (1 - self._momentum) * value
+        )
         self._denominator = self._momentum * self._denominator + (1 - self._momentum)
 
     @property
@@ -85,7 +93,9 @@ def search_func(
     arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter()
     end = time.time()
     network.train()
-    for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(xloader):
+    for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(
+        xloader
+    ):
         scheduler.update(None, 1.0 * step / len(xloader))
         base_inputs = base_inputs.cuda(non_blocking=True)
         arch_inputs = arch_inputs.cuda(non_blocking=True)
@@ -101,7 +111,9 @@ def search_func(
         base_loss.backward()
         w_optimizer.step()
         # record
-        base_prec1, base_prec5 = obtain_accuracy(logits.data, base_targets.data, topk=(1, 5))
+        base_prec1, base_prec5 = obtain_accuracy(
+            logits.data, base_targets.data, topk=(1, 5)
+        )
         base_losses.update(base_loss.item(), base_inputs.size(0))
         base_top1.update(base_prec1.item(), base_inputs.size(0))
         base_top5.update(base_prec5.item(), base_inputs.size(0))
@@ -110,7 +122,9 @@ def search_func(
         network.zero_grad()
         a_optimizer.zero_grad()
         _, logits, log_probs = network(arch_inputs)
-        arch_prec1, arch_prec5 = obtain_accuracy(logits.data, arch_targets.data, topk=(1, 5))
+        arch_prec1, arch_prec5 = obtain_accuracy(
+            logits.data, arch_targets.data, topk=(1, 5)
+        )
         if algo == "mask_rl":
             with torch.no_grad():
                 RL_BASELINE_EMA.update(arch_prec1.item())
@@ -134,7 +148,11 @@ def search_func(
         end = time.time()
 
         if step % print_freq == 0 or step + 1 == len(xloader):
-            Sstr = "*SEARCH* " + time_string() + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader))
+            Sstr = (
+                "*SEARCH* "
+                + time_string()
+                + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader))
+            )
             Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
                 batch_time=batch_time, data_time=data_time
             )
@@ -145,7 +163,14 @@ def search_func(
                 loss=arch_losses, top1=arch_top1, top5=arch_top5
             )
             logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Astr)
-    return base_losses.avg, base_top1.avg, base_top5.avg, arch_losses.avg, arch_top1.avg, arch_top5.avg
+    return (
+        base_losses.avg,
+        base_top1.avg,
+        base_top5.avg,
+        arch_losses.avg,
+        arch_top1.avg,
+        arch_top5.avg,
+    )
 
 
 def valid_func(xloader, network, criterion, logger):
@@ -162,7 +187,9 @@ def valid_func(xloader, network, criterion, logger):
             _, logits, _ = network(arch_inputs.cuda(non_blocking=True))
             arch_loss = criterion(logits, arch_targets)
             # record
-            arch_prec1, arch_prec5 = obtain_accuracy(logits.data, arch_targets.data, topk=(1, 5))
+            arch_prec1, arch_prec5 = obtain_accuracy(
+                logits.data, arch_targets.data, topk=(1, 5)
+            )
             arch_losses.update(arch_loss.item(), arch_inputs.size(0))
             arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
             arch_top5.update(arch_prec5.item(), arch_inputs.size(0))
@@ -181,11 +208,17 @@ def main(xargs):
     prepare_seed(xargs.rand_seed)
     logger = prepare_logger(args)
 
-    train_data, valid_data, xshape, class_num = get_datasets(xargs.dataset, xargs.data_path, -1)
+    train_data, valid_data, xshape, class_num = get_datasets(
+        xargs.dataset, xargs.data_path, -1
+    )
     if xargs.overwite_epochs is None:
         extra_info = {"class_num": class_num, "xshape": xshape}
     else:
-        extra_info = {"class_num": class_num, "xshape": xshape, "epochs": xargs.overwite_epochs}
+        extra_info = {
+            "class_num": class_num,
+            "xshape": xshape,
+            "epochs": xargs.overwite_epochs,
+        }
     config = load_config(xargs.config_path, extra_info, logger)
     search_loader, train_loader, valid_loader = get_nas_search_loaders(
         train_data,
@@ -223,7 +256,9 @@ def main(xargs):
     search_model.set_algo(xargs.algo)
     logger.log("{:}".format(search_model))
 
-    w_optimizer, w_scheduler, criterion = get_optim_scheduler(search_model.weights, config)
+    w_optimizer, w_scheduler, criterion = get_optim_scheduler(
+        search_model.weights, config
+    )
     a_optimizer = torch.optim.Adam(
         search_model.alphas,
         lr=xargs.arch_learning_rate,
@@ -244,13 +279,23 @@ def main(xargs):
         api = None
     logger.log("{:} create API = {:} done".format(time_string(), api))
 
-    last_info, model_base_path, model_best_path = logger.path("info"), logger.path("model"), logger.path("best")
+    last_info, model_base_path, model_best_path = (
+        logger.path("info"),
+        logger.path("model"),
+        logger.path("best"),
+    )
     network, criterion = search_model.cuda(), criterion.cuda()  # use a single GPU
 
-    last_info, model_base_path, model_best_path = logger.path("info"), logger.path("model"), logger.path("best")
+    last_info, model_base_path, model_best_path = (
+        logger.path("info"),
+        logger.path("model"),
+        logger.path("best"),
+    )
 
     if last_info.exists():  # automatically resume from previous checkpoint
-        logger.log("=> loading checkpoint of the last-info '{:}' start".format(last_info))
+        logger.log(
+            "=> loading checkpoint of the last-info '{:}' start".format(last_info)
+        )
         last_info = torch.load(last_info)
         start_epoch = last_info["epoch"]
         checkpoint = torch.load(last_info["last_checkpoint"])
@@ -261,7 +306,9 @@ def main(xargs):
         w_optimizer.load_state_dict(checkpoint["w_optimizer"])
         a_optimizer.load_state_dict(checkpoint["a_optimizer"])
         logger.log(
-            "=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(last_info, start_epoch)
+            "=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(
+                last_info, start_epoch
+            )
         )
     else:
         logger.log("=> do not find the last-info file : {:}".format(last_info))
@@ -276,26 +323,47 @@ def main(xargs):
     )
     for epoch in range(start_epoch, total_epoch):
         w_scheduler.update(epoch, 0.0)
-        need_time = "Time Left: {:}".format(convert_secs2time(epoch_time.val * (total_epoch - epoch), True))
+        need_time = "Time Left: {:}".format(
+            convert_secs2time(epoch_time.val * (total_epoch - epoch), True)
+        )
         epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
 
-        if xargs.warmup_ratio is None or xargs.warmup_ratio <= float(epoch) / total_epoch:
+        if (
+            xargs.warmup_ratio is None
+            or xargs.warmup_ratio <= float(epoch) / total_epoch
+        ):
             enable_controller = True
             network.set_warmup_ratio(None)
         else:
             enable_controller = False
-            network.set_warmup_ratio(1.0 - float(epoch) / total_epoch / xargs.warmup_ratio)
+            network.set_warmup_ratio(
+                1.0 - float(epoch) / total_epoch / xargs.warmup_ratio
+            )
 
         logger.log(
             "\n[Search the {:}-th epoch] {:}, LR={:}, controller-warmup={:}, enable_controller={:}".format(
-                epoch_str, need_time, min(w_scheduler.get_lr()), network.warmup_ratio, enable_controller
+                epoch_str,
+                need_time,
+                min(w_scheduler.get_lr()),
+                network.warmup_ratio,
+                enable_controller,
             )
         )
 
         if xargs.algo == "mask_gumbel" or xargs.algo == "tas":
-            network.set_tau(xargs.tau_max - (xargs.tau_max - xargs.tau_min) * epoch / (total_epoch - 1))
+            network.set_tau(
+                xargs.tau_max
+                - (xargs.tau_max - xargs.tau_min) * epoch / (total_epoch - 1)
+            )
             logger.log("[RESET tau as : {:}]".format(network.tau))
-        search_w_loss, search_w_top1, search_w_top5, search_a_loss, search_a_top1, search_a_top5 = search_func(
+        (
+            search_w_loss,
+            search_w_top1,
+            search_w_top5,
+            search_a_loss,
+            search_a_top1,
+            search_a_top5,
+        ) = search_func(
             search_loader,
             network,
             criterion,
@@ -322,7 +390,9 @@ def main(xargs):
 
         genotype = network.genotype
         logger.log("[{:}] - [get_best_arch] : {:}".format(epoch_str, genotype))
-        valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(valid_loader, network, criterion, logger)
+        valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
+            valid_loader, network, criterion, logger
+        )
         logger.log(
             "[{:}] evaluate : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}% | {:}".format(
                 epoch_str, valid_a_loss, valid_a_top1, valid_a_top5, genotype
@@ -331,7 +401,9 @@ def main(xargs):
         valid_accuracies[epoch] = valid_a_top1
 
         genotypes[epoch] = genotype
-        logger.log("<<<--->>> The {:}-th epoch : {:}".format(epoch_str, genotypes[epoch]))
+        logger.log(
+            "<<<--->>> The {:}-th epoch : {:}".format(epoch_str, genotypes[epoch])
+        )
         # save checkpoint
         save_path = save_checkpoint(
             {
@@ -369,8 +441,14 @@ def main(xargs):
     genotype = network.genotype
     search_time.update(time.time() - start_time)
 
-    valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(valid_loader, network, criterion, logger)
-    logger.log("Last : the gentotype is : {:}, with the validation accuracy of {:.3f}%.".format(genotype, valid_a_top1))
+    valid_a_loss, valid_a_top1, valid_a_top5 = valid_func(
+        valid_loader, network, criterion, logger
+    )
+    logger.log(
+        "Last : the gentotype is : {:}, with the validation accuracy of {:.3f}%.".format(
+            genotype, valid_a_top1
+        )
+    )
 
     logger.log("\n" + "-" * 100)
     # check the performance from the architecture dataset
@@ -393,8 +471,19 @@ if __name__ == "__main__":
         choices=["cifar10", "cifar100", "ImageNet16-120"],
         help="Choose between Cifar10/100 and ImageNet-16.",
     )
-    parser.add_argument("--search_space", type=str, default="sss", choices=["sss"], help="The search space name.")
-    parser.add_argument("--algo", type=str, choices=["tas", "mask_gumbel", "mask_rl"], help="The search space name.")
+    parser.add_argument(
+        "--search_space",
+        type=str,
+        default="sss",
+        choices=["sss"],
+        help="The search space name.",
+    )
+    parser.add_argument(
+        "--algo",
+        type=str,
+        choices=["tas", "mask_gumbel", "mask_rl"],
+        help="The search space name.",
+    )
     parser.add_argument(
         "--genotype",
         type=str,
@@ -402,13 +491,23 @@ if __name__ == "__main__":
         help="The genotype.",
     )
     parser.add_argument(
-        "--use_api", type=int, default=1, choices=[0, 1], help="Whether use API or not (which will cost much memory)."
+        "--use_api",
+        type=int,
+        default=1,
+        choices=[0, 1],
+        help="Whether use API or not (which will cost much memory).",
     )
     # FOR GDAS
-    parser.add_argument("--tau_min", type=float, default=0.1, help="The minimum tau for Gumbel Softmax.")
-    parser.add_argument("--tau_max", type=float, default=10, help="The maximum tau for Gumbel Softmax.")
+    parser.add_argument(
+        "--tau_min", type=float, default=0.1, help="The minimum tau for Gumbel Softmax."
+    )
+    parser.add_argument(
+        "--tau_max", type=float, default=10, help="The maximum tau for Gumbel Softmax."
+    )
     # FOR ALL
-    parser.add_argument("--warmup_ratio", type=float, help="The warmup ratio, if None, not use warmup.")
+    parser.add_argument(
+        "--warmup_ratio", type=float, help="The warmup ratio, if None, not use warmup."
+    )
     #
     parser.add_argument(
         "--track_running_stats",
@@ -418,7 +517,11 @@ if __name__ == "__main__":
         help="Whether use track_running_stats or not in the BN layer.",
     )
     parser.add_argument(
-        "--affine", type=int, default=0, choices=[0, 1], help="Whether use affine=True or False in the BN layer."
+        "--affine",
+        type=int,
+        default=0,
+        choices=[0, 1],
+        help="Whether use affine=True or False in the BN layer.",
     )
     parser.add_argument(
         "--config_path",
@@ -427,25 +530,57 @@ if __name__ == "__main__":
         help="The path of configuration.",
     )
     parser.add_argument(
-        "--overwite_epochs", type=int, help="The number of epochs to overwrite that value in config files."
+        "--overwite_epochs",
+        type=int,
+        help="The number of epochs to overwrite that value in config files.",
     )
     # architecture leraning rate
-    parser.add_argument("--arch_learning_rate", type=float, default=3e-4, help="learning rate for arch encoding")
-    parser.add_argument("--arch_weight_decay", type=float, default=1e-3, help="weight decay for arch encoding")
-    parser.add_argument("--arch_eps", type=float, default=1e-8, help="weight decay for arch encoding")
+    parser.add_argument(
+        "--arch_learning_rate",
+        type=float,
+        default=3e-4,
+        help="learning rate for arch encoding",
+    )
+    parser.add_argument(
+        "--arch_weight_decay",
+        type=float,
+        default=1e-3,
+        help="weight decay for arch encoding",
+    )
+    parser.add_argument(
+        "--arch_eps", type=float, default=1e-8, help="weight decay for arch encoding"
+    )
     # log
-    parser.add_argument("--workers", type=int, default=2, help="number of data loading workers (default: 2)")
-    parser.add_argument("--save_dir", type=str, default="./output/search", help="Folder to save checkpoints and log.")
-    parser.add_argument("--print_freq", type=int, default=200, help="print frequency (default: 200)")
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=2,
+        help="number of data loading workers (default: 2)",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="./output/search",
+        help="Folder to save checkpoints and log.",
+    )
+    parser.add_argument(
+        "--print_freq", type=int, default=200, help="print frequency (default: 200)"
+    )
     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{:}-AWD{:}-WARM{:}".format(
-        args.algo, args.affine, args.track_running_stats, args.arch_weight_decay, args.warmup_ratio
+        args.algo,
+        args.affine,
+        args.track_running_stats,
+        args.arch_weight_decay,
+        args.warmup_ratio,
     )
     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)
+    args.save_dir = os.path.join(
+        "{:}-{:}".format(args.save_dir, args.search_space), args.dataset, dirname
+    )
 
     main(args)