From 46b92e37e2ea3d735f6d682293dd7ed21baa65cf Mon Sep 17 00:00:00 2001 From: D-X-Y <280835372@qq.com> Date: Tue, 1 Dec 2020 22:25:23 +0800 Subject: [PATCH] Add get_torch_home func for NATS-Bench --- exps/NATS-Bench/draw-fig2_5.py | 2 +- exps/NATS-Bench/draw-fig8.py | 175 ++++++++++++++++++++++++++++++ exps/NATS-Bench/draw-ranks.py | 96 ++++++++++++++++ exps/NATS-algos/regularized_ea.py | 10 +- lib/nats_bench/api_size.py | 5 +- lib/nats_bench/api_topology.py | 5 +- lib/nats_bench/api_utils.py | 11 ++ 7 files changed, 294 insertions(+), 10 deletions(-) create mode 100644 exps/NATS-Bench/draw-fig8.py create mode 100644 exps/NATS-Bench/draw-ranks.py diff --git a/exps/NATS-Bench/draw-fig2_5.py b/exps/NATS-Bench/draw-fig2_5.py index 44d563b..f5b44a8 100644 --- a/exps/NATS-Bench/draw-fig2_5.py +++ b/exps/NATS-Bench/draw-fig2_5.py @@ -385,7 +385,7 @@ def visualize_all_rank_info(api, vis_save_dir, indicator): if __name__ == '__main__': - parser = argparse.ArgumentParser(description='NAS-Bench-X', formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser = argparse.ArgumentParser(description='NATS-Bench', formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--save_dir', type=str, default='output/vis-nas-bench', help='Folder to save checkpoints and log.') # use for train the model args = parser.parse_args() diff --git a/exps/NATS-Bench/draw-fig8.py b/exps/NATS-Bench/draw-fig8.py new file mode 100644 index 0000000..c3a5b72 --- /dev/null +++ b/exps/NATS-Bench/draw-fig8.py @@ -0,0 +1,175 @@ +############################################################### +# NATS-Bench (https://arxiv.org/pdf/2009.00437.pdf) # +# The code to draw Figure 6 in our paper. # +############################################################### +# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.06 # +############################################################### +# Usage: python exps/NATS-Bench/draw-fig8.py # +############################################################### +import os, gc, sys, time, torch, argparse +import numpy as np +from typing import List, Text, Dict, Any +from shutil import copyfile +from collections import defaultdict, OrderedDict +from copy import deepcopy +from pathlib import Path +import matplotlib +import seaborn as sns +matplotlib.use('agg') +import matplotlib.pyplot as plt +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 nats_bench import create +from log_utils import time_string + +plt.rcParams.update({ + "text.usetex": True, + "font.family": "sans-serif", + "font.sans-serif": ["Helvetica"]}) +## for Palatino and other serif fonts use: +plt.rcParams.update({ + "text.usetex": True, + "font.family": "serif", + "font.serif": ["Palatino"], +}) + + +def fetch_data(root_dir='./output/search', search_space='tss', dataset=None): + ss_dir = '{:}-{:}'.format(root_dir, search_space) + alg2all = OrderedDict() + # alg2name['REINFORCE'] = 'REINFORCE-0.01' + # alg2name['RANDOM'] = 'RANDOM' + # alg2name['BOHB'] = 'BOHB' + if dataset == 'cifar10': + suffixes = ['-T200000', '-T200000-FULL'] + elif dataset == 'cifar100': + suffixes = ['-T40000', '-T40000-FULL'] + elif search_space == 'tss': + suffixes = ['-T120000', '-T120000-FULL'] + elif search_space == 'sss': + suffixes = ['-T60000', '-T60000-FULL'] + else: + raise ValueError('Unkonwn dataset : {:}'.format(dataset)) + if search_space == 'tss': + hp = '$\mathcal{H}^{1}$' + elif search_space == 'sss': + hp = '$\mathcal{H}^{2}$' + else: + raise ValueError('Unkonwn search space: {:}'.format(search_space)) + + alg2all[r'REA ($\mathcal{H}^{0}$)'] = dict( + path=os.path.join(ss_dir, dataset + suffixes[0], 'R-EA-SS3', 'results.pth'), + color='b', linestyle='-') + alg2all[r'REA ({:})'.format(hp)] = dict( + path=os.path.join(ss_dir, dataset + suffixes[1], 'R-EA-SS3', 'results.pth'), + color='b', linestyle='--') + + for alg, xdata in alg2all.items(): + data = torch.load(xdata['path']) + for index, info in data.items(): + info['time_w_arch'] = [(x, y) for x, y in zip(info['all_total_times'], info['all_archs'])] + for j, arch in enumerate(info['all_archs']): + assert arch != -1, 'invalid arch from {:} {:} {:} ({:}, {:})'.format(alg, search_space, dataset, index, j) + xdata['data'] = data + return alg2all + + +def query_performance(api, data, dataset, ticket): + 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_size_space 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) + return np.mean(results), np.std(results) + + +y_min_s = {('cifar10', 'tss'): 90, + ('cifar10', 'sss'): 90, + ('cifar100', 'tss'): 65, + ('cifar100', 'sss'): 65, + ('ImageNet16-120', 'tss'): 36, + ('ImageNet16-120', 'sss'): 40} + +y_max_s = {('cifar10', 'tss'): 94.5, + ('cifar10', 'sss'): 94.5, + ('cifar100', 'tss'): 72.5, + ('cifar100', 'sss'): 70.5, + ('ImageNet16-120', 'tss'): 46, + ('ImageNet16-120', 'sss'): 46} + +x_axis_s = {('cifar10', 'tss'): 200000, + ('cifar10', 'sss'): 200000, + ('cifar100', 'tss'): 400, + ('cifar100', 'sss'): 400, + ('ImageNet16-120', 'tss'): 1200, + ('ImageNet16-120', 'sss'): 600} + +name2label = {'cifar10': 'CIFAR-10', + 'cifar100': 'CIFAR-100', + 'ImageNet16-120': 'ImageNet-16-120'} + +spaces2latex = {'tss': r'$\mathcal{S}_{t}$', + 'sss': r'$\mathcal{S}_{s}$',} + +def visualize_curve(api_dict, vis_save_dir): + vis_save_dir = vis_save_dir.resolve() + vis_save_dir.mkdir(parents=True, exist_ok=True) + + dpi, width, height = 250, 4000, 2400 + figsize = width / float(dpi), height / float(dpi) + LabelSize, LegendFontsize = 16, 16 + + def sub_plot_fn(ax, search_space, dataset): + max_time = x_axis_s[(dataset, search_space)] + alg2data = fetch_data(search_space=search_space, dataset=dataset) + alg2accuracies = OrderedDict() + total_tickets = 200 + time_tickets = [float(i) / total_tickets * int(max_time) for i in range(total_tickets)] + ax.set_xlim(0, x_axis_s[(dataset, search_space)]) + ax.set_ylim(y_min_s[(dataset, search_space)], + y_max_s[(dataset, search_space)]) + for idx, (alg, xdata) in enumerate(alg2data.items()): + accuracies = [] + for ticket in time_tickets: + # import pdb; pdb.set_trace() + accuracy, accuracy_std = query_performance( + api_dict[search_space], xdata['data'], dataset, ticket) + accuracies.append(accuracy) + # print('{:} plot alg : {:10s}, final accuracy = {:.2f}$\pm${:.2f}'.format(time_string(), alg, accuracy, accuracy_std)) + print('{:} plot alg : {:10s} on {:}'.format(time_string(), alg, search_space)) + alg2accuracies[alg] = accuracies + ax.plot(time_tickets, accuracies, c=xdata['color'], linestyle=xdata['linestyle'], label='{:}'.format(alg)) + ax.set_xlabel('Estimated wall-clock time', fontsize=LabelSize) + ax.set_ylabel('Test accuracy', fontsize=LabelSize) + ax.set_title(r'Searching results on {:} for {:}'.format(name2label[dataset], spaces2latex[search_space]), + fontsize=LabelSize+4) + ax.legend(loc=4, fontsize=LegendFontsize) + + fig, axs = plt.subplots(1, 2, figsize=figsize) + sub_plot_fn(axs[0], 'tss', 'cifar10') + sub_plot_fn(axs[1], 'sss', 'cifar10') + save_path = (vis_save_dir / 'full-curve.png').resolve() + fig.savefig(save_path, dpi=dpi, bbox_inches='tight', format='png') + print ('{:} save into {:}'.format(time_string(), save_path)) + plt.close('all') + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size', formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument('--save_dir', type=str, default='output/vis-nas-bench/nas-algos-vs-h', help='Folder to save checkpoints and log.') + args = parser.parse_args() + + save_dir = Path(args.save_dir) + + api_tss = create(None, 'tss', fast_mode=True, verbose=False) + api_sss = create(None, 'sss', fast_mode=True, verbose=False) + visualize_curve(dict(tss=api_tss, sss=api_sss), save_dir) diff --git a/exps/NATS-Bench/draw-ranks.py b/exps/NATS-Bench/draw-ranks.py new file mode 100644 index 0000000..e8f7d65 --- /dev/null +++ b/exps/NATS-Bench/draw-ranks.py @@ -0,0 +1,96 @@ +############################################################### +# NATS-Bench (https://arxiv.org/pdf/2009.00437.pdf) # +# The code to draw Figure 2 / 3 / 4 / 5 in our paper. # +############################################################### +# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.06 # +############################################################### +# Usage: python exps/NATS-Bench/draw-ranks.py # +############################################################### +import os, sys, time, torch, argparse +import scipy +import numpy as np +from typing import List, Text, Dict, Any +from shutil import copyfile +from collections import defaultdict +from copy import deepcopy +from pathlib import Path +import matplotlib +import seaborn as sns +matplotlib.use('agg') +import matplotlib.pyplot as plt +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 log_utils import time_string +from models import get_cell_based_tiny_net +from nats_bench import create + + +def visualize_relative_info(api, vis_save_dir, indicator): + vis_save_dir = vis_save_dir.resolve() + # print ('{:} start to visualize {:} information'.format(time_string(), api)) + vis_save_dir.mkdir(parents=True, exist_ok=True) + + cifar010_cache_path = vis_save_dir / '{:}-cache-{:}-info.pth'.format('cifar10', indicator) + cifar100_cache_path = vis_save_dir / '{:}-cache-{:}-info.pth'.format('cifar100', indicator) + imagenet_cache_path = vis_save_dir / '{:}-cache-{:}-info.pth'.format('ImageNet16-120', indicator) + cifar010_info = torch.load(cifar010_cache_path) + cifar100_info = torch.load(cifar100_cache_path) + imagenet_info = torch.load(imagenet_cache_path) + indexes = list(range(len(cifar010_info['params']))) + + print ('{:} start to visualize relative ranking'.format(time_string())) + + cifar010_ord_indexes = sorted(indexes, key=lambda i: cifar010_info['test_accs'][i]) + cifar100_ord_indexes = sorted(indexes, key=lambda i: cifar100_info['test_accs'][i]) + imagenet_ord_indexes = sorted(indexes, key=lambda i: imagenet_info['test_accs'][i]) + + cifar100_labels, imagenet_labels = [], [] + for idx in cifar010_ord_indexes: + cifar100_labels.append( cifar100_ord_indexes.index(idx) ) + imagenet_labels.append( imagenet_ord_indexes.index(idx) ) + print ('{:} prepare data done.'.format(time_string())) + + dpi, width, height = 200, 1400, 800 + figsize = width / float(dpi), height / float(dpi) + LabelSize, LegendFontsize = 18, 12 + resnet_scale, resnet_alpha = 120, 0.5 + + fig = plt.figure(figsize=figsize) + ax = fig.add_subplot(111) + plt.xlim(min(indexes), max(indexes)) + plt.ylim(min(indexes), max(indexes)) + # plt.ylabel('y').set_rotation(30) + plt.yticks(np.arange(min(indexes), max(indexes), max(indexes)//3), fontsize=LegendFontsize, rotation='vertical') + plt.xticks(np.arange(min(indexes), max(indexes), max(indexes)//5), fontsize=LegendFontsize) + ax.scatter(indexes, cifar100_labels, marker='^', s=0.5, c='tab:green', alpha=0.8) + ax.scatter(indexes, imagenet_labels, marker='*', s=0.5, c='tab:red' , alpha=0.8) + ax.scatter(indexes, indexes , marker='o', s=0.5, c='tab:blue' , alpha=0.8) + ax.scatter([-1], [-1], marker='o', s=100, c='tab:blue' , label='CIFAR-10') + ax.scatter([-1], [-1], marker='^', s=100, c='tab:green', label='CIFAR-100') + ax.scatter([-1], [-1], marker='*', s=100, c='tab:red' , label='ImageNet-16-120') + plt.grid(zorder=0) + ax.set_axisbelow(True) + plt.legend(loc=0, fontsize=LegendFontsize) + ax.set_xlabel('architecture ranking in CIFAR-10', fontsize=LabelSize) + ax.set_ylabel('architecture ranking', fontsize=LabelSize) + save_path = (vis_save_dir / '{:}-relative-rank.pdf'.format(indicator)).resolve() + fig.savefig(save_path, dpi=dpi, bbox_inches='tight', format='pdf') + save_path = (vis_save_dir / '{:}-relative-rank.png'.format(indicator)).resolve() + fig.savefig(save_path, dpi=dpi, bbox_inches='tight', format='png') + print ('{:} save into {:}'.format(time_string(), save_path)) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='NATS-Bench', formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument('--save_dir', type=str, default='output/vis-nas-bench/rank-stability', help='Folder to save checkpoints and log.') + # use for train the model + args = parser.parse_args() + + to_save_dir = Path(args.save_dir) + + # Figure 2 + visualize_relative_info(None, to_save_dir, 'tss') + visualize_relative_info(None, to_save_dir, 'sss') \ No newline at end of file diff --git a/exps/NATS-algos/regularized_ea.py b/exps/NATS-algos/regularized_ea.py index 14ba2fd..be63992 100644 --- a/exps/NATS-algos/regularized_ea.py +++ b/exps/NATS-algos/regularized_ea.py @@ -9,6 +9,7 @@ # python ./exps/NATS-algos/regularized_ea.py --dataset cifar10 --search_space sss --ea_cycles 200 --ea_population 10 --ea_sample_size 3 --rand_seed 1 # python ./exps/NATS-algos/regularized_ea.py --dataset cifar100 --search_space sss --ea_cycles 200 --ea_population 10 --ea_sample_size 3 --rand_seed 1 # python ./exps/NATS-algos/regularized_ea.py --dataset ImageNet16-120 --search_space sss --ea_cycles 200 --ea_population 10 --ea_sample_size 3 --rand_seed 1 +# python ./exps/NATS-algos/regularized_ea.py --dataset ${dataset} --search_space ${search_space} --time_budget ${time_budget} --ea_cycles 200 --ea_population 10 --ea_sample_size 3 --use_proxy 0 ################################################################## import os, sys, time, glob, random, argparse import numpy as np, collections @@ -119,10 +120,8 @@ def regularized_evolution(cycles, population_size, sample_size, time_budget, ran while len(population) < population_size: model = Model() model.arch = random_arch() - if use_proxy: - model.accuracy, _, _, total_cost = api.simulate_train_eval(model.arch, dataset, hp='12') - else: - model.accuracy, _, _, total_cost = api.simulate_train_eval(model.arch, dataset, hp=api.full_train_epochs) + model.accuracy, _, _, total_cost = api.simulate_train_eval( + model.arch, dataset, hp='12' if use_proxy else api.full_train_epochs) # Append the info population.append(model) history.append((model.accuracy, model.arch)) @@ -146,7 +145,8 @@ def regularized_evolution(cycles, population_size, sample_size, time_budget, ran # Create the child model and store it. child = Model() child.arch = mutate_arch(parent.arch) - child.accuracy, _, _, total_cost = api.simulate_train_eval(child.arch, dataset, hp='12') + child.accuracy, _, _, total_cost = api.simulate_train_eval( + child.arch, dataset, hp='12' if use_proxy else api.full_train_epochs) # Append the info population.append(child) history.append((child.accuracy, child.arch)) diff --git a/lib/nats_bench/api_size.py b/lib/nats_bench/api_size.py index 6eab753..b6717c3 100644 --- a/lib/nats_bench/api_size.py +++ b/lib/nats_bench/api_size.py @@ -17,6 +17,7 @@ from typing import Dict, Optional, Text, Union, Any from nats_bench.api_utils import ArchResults from nats_bench.api_utils import NASBenchMetaAPI +from nats_bench.api_utils import get_torch_home from nats_bench.api_utils import nats_is_dir from nats_bench.api_utils import nats_is_file from nats_bench.api_utils import PICKLE_EXT @@ -88,10 +89,10 @@ class NATSsize(NASBenchMetaAPI): if file_path_or_dict is None: if self._fast_mode: self._archive_dir = os.path.join( - os.environ['TORCH_HOME'], '{:}-simple'.format(ALL_BASE_NAMES[-1])) + get_torch_home(), '{:}-simple'.format(ALL_BASE_NAMES[-1])) else: file_path_or_dict = os.path.join( - os.environ['TORCH_HOME'], '{:}.{:}'.format( + get_torch_home(), '{:}.{:}'.format( ALL_BASE_NAMES[-1], PICKLE_EXT)) print('{:} Try to use the default NATS-Bench (size) path from ' 'fast_mode={:} and path={:}.'.format(time_string(), self._fast_mode, diff --git a/lib/nats_bench/api_topology.py b/lib/nats_bench/api_topology.py index f4211c4..fd2914c 100644 --- a/lib/nats_bench/api_topology.py +++ b/lib/nats_bench/api_topology.py @@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Text, Union from nats_bench.api_utils import ArchResults from nats_bench.api_utils import NASBenchMetaAPI +from nats_bench.api_utils import get_torch_home from nats_bench.api_utils import nats_is_dir from nats_bench.api_utils import nats_is_file from nats_bench.api_utils import PICKLE_EXT @@ -88,10 +89,10 @@ class NATStopology(NASBenchMetaAPI): if file_path_or_dict is None: if self._fast_mode: self._archive_dir = os.path.join( - os.environ['TORCH_HOME'], '{:}-simple'.format(ALL_BASE_NAMES[-1])) + get_torch_home(), '{:}-simple'.format(ALL_BASE_NAMES[-1])) else: file_path_or_dict = os.path.join( - os.environ['TORCH_HOME'], '{:}.{:}'.format( + get_torch_home(), '{:}.{:}'.format( ALL_BASE_NAMES[-1], PICKLE_EXT)) print('{:} Try to use the default NATS-Bench (topology) path from ' 'fast_mode={:} and path={:}.'.format(time_string(), self._fast_mode, file_path_or_dict)) diff --git a/lib/nats_bench/api_utils.py b/lib/nats_bench/api_utils.py index 86a7a02..f858822 100644 --- a/lib/nats_bench/api_utils.py +++ b/lib/nats_bench/api_utils.py @@ -45,6 +45,17 @@ def get_file_system(): return _FILE_SYSTEM +def get_torch_home(): + if 'TORCH_HOME' in os.environ: + return os.environ['TORCH_HOME'] + elif 'HOME' in os.environ: + return os.path.join(os.environ['HOME'], '.torch') + else: + raise ValueError('Did not find HOME in os.environ. ' + 'Please at least setup the path of HOME or TORCH_HOME ' + 'in the environment.') + + def nats_is_dir(file_path): if _FILE_SYSTEM == 'default': return os.path.isdir(file_path)