189 lines
7.1 KiB
Python
189 lines
7.1 KiB
Python
###############################################################
|
|
# NATS-Bench (arxiv.org/pdf/2009.00437.pdf), IEEE TPAMI 2021 #
|
|
# 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, 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 log_utils import time_string
|
|
from models import get_cell_based_tiny_net
|
|
from nats_bench import create
|
|
|
|
|
|
name2label = {
|
|
"cifar10": "CIFAR-10",
|
|
"cifar100": "CIFAR-100",
|
|
"ImageNet16-120": "ImageNet-16-120",
|
|
}
|
|
|
|
|
|
def visualize_relative_info(vis_save_dir, search_space, indicator, topk):
|
|
vis_save_dir = vis_save_dir.resolve()
|
|
print(
|
|
"{:} start to visualize {:} with top-{:} information".format(
|
|
time_string(), search_space, topk
|
|
)
|
|
)
|
|
vis_save_dir.mkdir(parents=True, exist_ok=True)
|
|
cache_file_path = vis_save_dir / "cache-{:}-info.pth".format(search_space)
|
|
datasets = ["cifar10", "cifar100", "ImageNet16-120"]
|
|
if not cache_file_path.exists():
|
|
api = create(None, search_space, fast_mode=False, verbose=False)
|
|
all_infos = OrderedDict()
|
|
for index in range(len(api)):
|
|
all_info = OrderedDict()
|
|
for dataset in datasets:
|
|
info_less = api.get_more_info(index, dataset, hp="12", is_random=False)
|
|
info_more = api.get_more_info(
|
|
index, dataset, hp=api.full_train_epochs, is_random=False
|
|
)
|
|
all_info[dataset] = dict(
|
|
less=info_less["test-accuracy"], more=info_more["test-accuracy"]
|
|
)
|
|
all_infos[index] = all_info
|
|
torch.save(all_infos, cache_file_path)
|
|
print("{:} save all cache data into {:}".format(time_string(), cache_file_path))
|
|
else:
|
|
api = create(None, search_space, fast_mode=True, verbose=False)
|
|
all_infos = torch.load(cache_file_path)
|
|
|
|
dpi, width, height = 250, 5000, 1300
|
|
figsize = width / float(dpi), height / float(dpi)
|
|
LabelSize, LegendFontsize = 16, 16
|
|
|
|
fig, axs = plt.subplots(1, 3, figsize=figsize)
|
|
datasets = ["cifar10", "cifar100", "ImageNet16-120"]
|
|
|
|
def sub_plot_fn(ax, dataset, indicator):
|
|
performances = []
|
|
# pickup top 10% architectures
|
|
for _index in range(len(api)):
|
|
performances.append((all_infos[_index][dataset][indicator], _index))
|
|
performances = sorted(performances, reverse=True)
|
|
performances = performances[: int(len(api) * topk * 0.01)]
|
|
selected_indexes = [x[1] for x in performances]
|
|
print(
|
|
"{:} plot {:10s} with {:}, {:} architectures".format(
|
|
time_string(), dataset, indicator, len(selected_indexes)
|
|
)
|
|
)
|
|
standard_scores = []
|
|
random_scores = []
|
|
for idx in selected_indexes:
|
|
standard_scores.append(
|
|
api.get_more_info(
|
|
idx,
|
|
dataset,
|
|
hp=api.full_train_epochs if indicator == "more" else "12",
|
|
is_random=False,
|
|
)["test-accuracy"]
|
|
)
|
|
random_scores.append(
|
|
api.get_more_info(
|
|
idx,
|
|
dataset,
|
|
hp=api.full_train_epochs if indicator == "more" else "12",
|
|
is_random=True,
|
|
)["test-accuracy"]
|
|
)
|
|
indexes = list(range(len(selected_indexes)))
|
|
standard_indexes = sorted(indexes, key=lambda i: standard_scores[i])
|
|
random_indexes = sorted(indexes, key=lambda i: random_scores[i])
|
|
random_labels = []
|
|
for idx in standard_indexes:
|
|
random_labels.append(random_indexes.index(idx))
|
|
for tick in ax.get_xticklabels():
|
|
tick.set_fontsize(LabelSize - 3)
|
|
for tick in ax.get_yticklabels():
|
|
tick.set_rotation(25)
|
|
tick.set_fontsize(LabelSize - 3)
|
|
ax.set_xlim(0, len(indexes))
|
|
ax.set_ylim(0, len(indexes))
|
|
ax.set_yticks(np.arange(min(indexes), max(indexes), max(indexes) // 3))
|
|
ax.set_xticks(np.arange(min(indexes), max(indexes), max(indexes) // 5))
|
|
ax.scatter(indexes, random_labels, marker="^", s=0.5, c="tab:green", 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="Average Over Multi-Trials",
|
|
)
|
|
ax.scatter(
|
|
[-1],
|
|
[-1],
|
|
marker="^",
|
|
s=100,
|
|
c="tab:green",
|
|
label="Randomly Selected Trial",
|
|
)
|
|
|
|
coef, p = scipy.stats.kendalltau(standard_scores, random_scores)
|
|
ax.set_xlabel(
|
|
"architecture ranking in {:}".format(name2label[dataset]),
|
|
fontsize=LabelSize,
|
|
)
|
|
if dataset == "cifar10":
|
|
ax.set_ylabel("architecture ranking", fontsize=LabelSize)
|
|
ax.legend(loc=4, fontsize=LegendFontsize)
|
|
return coef
|
|
|
|
for dataset, ax in zip(datasets, axs):
|
|
rank_coef = sub_plot_fn(ax, dataset, indicator)
|
|
print(
|
|
"sub-plot {:} on {:} done, the ranking coefficient is {:.4f}.".format(
|
|
dataset, search_space, rank_coef
|
|
)
|
|
)
|
|
|
|
save_path = (
|
|
vis_save_dir / "{:}-rank-{:}-top{:}.pdf".format(search_space, indicator, topk)
|
|
).resolve()
|
|
fig.savefig(save_path, dpi=dpi, bbox_inches="tight", format="pdf")
|
|
save_path = (
|
|
vis_save_dir / "{:}-rank-{:}-top{:}.png".format(search_space, indicator, topk)
|
|
).resolve()
|
|
fig.savefig(save_path, dpi=dpi, bbox_inches="tight", format="png")
|
|
print("Save into {:}".format(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.",
|
|
)
|
|
args = parser.parse_args()
|
|
to_save_dir = Path(args.save_dir)
|
|
|
|
for topk in [1, 5, 10, 20]:
|
|
visualize_relative_info(to_save_dir, "tss", "more", topk)
|
|
visualize_relative_info(to_save_dir, "sss", "less", topk)
|
|
print("{:} : complete running this file : {:}".format(time_string(), __file__))
|