254 lines
9.4 KiB
Python
254 lines
9.4 KiB
Python
###############################################################
|
|
# NATS-Bench (arxiv.org/pdf/2009.00437.pdf), IEEE TPAMI 2021 #
|
|
# The code to draw Figure 7 in our paper. #
|
|
###############################################################
|
|
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.06 #
|
|
###############################################################
|
|
# Usage: python exps/NATS-Bench/draw-fig7.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
|
|
|
|
|
|
def get_valid_test_acc(api, arch, dataset):
|
|
is_size_space = api.search_space_name == "size"
|
|
if dataset == "cifar10":
|
|
xinfo = api.get_more_info(
|
|
arch, dataset=dataset, hp=90 if is_size_space else 200, is_random=False
|
|
)
|
|
test_acc = xinfo["test-accuracy"]
|
|
xinfo = api.get_more_info(
|
|
arch,
|
|
dataset="cifar10-valid",
|
|
hp=90 if is_size_space else 200,
|
|
is_random=False,
|
|
)
|
|
valid_acc = xinfo["valid-accuracy"]
|
|
else:
|
|
xinfo = api.get_more_info(
|
|
arch, dataset=dataset, hp=90 if is_size_space else 200, is_random=False
|
|
)
|
|
valid_acc = xinfo["valid-accuracy"]
|
|
test_acc = xinfo["test-accuracy"]
|
|
return (
|
|
valid_acc,
|
|
test_acc,
|
|
"validation = {:.2f}, test = {:.2f}\n".format(valid_acc, test_acc),
|
|
)
|
|
|
|
|
|
def fetch_data(
|
|
root_dir="./output/search", search_space="tss", dataset=None, suffix="-WARM0.3"
|
|
):
|
|
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":
|
|
alg2name["GDAS"] = "gdas-affine0_BN0-None"
|
|
alg2name["RSPS"] = "random-affine0_BN0-None"
|
|
alg2name["DARTS (1st)"] = "darts-v1-affine0_BN0-None"
|
|
alg2name["DARTS (2nd)"] = "darts-v2-affine0_BN0-None"
|
|
alg2name["ENAS"] = "enas-affine0_BN0-None"
|
|
alg2name["SETN"] = "setn-affine0_BN0-None"
|
|
else:
|
|
alg2name["channel-wise interpolation"] = "tas-affine0_BN0-AWD0.001{:}".format(
|
|
suffix
|
|
)
|
|
alg2name[
|
|
"masking + Gumbel-Softmax"
|
|
] = "mask_gumbel-affine0_BN0-AWD0.001{:}".format(suffix)
|
|
alg2name["masking + sampling"] = "mask_rl-affine0_BN0-AWD0.0{:}".format(suffix)
|
|
for alg, name in alg2name.items():
|
|
alg2path[alg] = os.path.join(ss_dir, dataset, name, "seed-{:}-last-info.pth")
|
|
alg2data = OrderedDict()
|
|
for alg, path in alg2path.items():
|
|
alg2data[alg], ok_num = [], 0
|
|
for seed in seeds:
|
|
xpath = path.format(seed)
|
|
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"))
|
|
try:
|
|
data = torch.load(
|
|
data["last_checkpoint"], map_location=torch.device("cpu")
|
|
)
|
|
except:
|
|
xpath = str(data["last_checkpoint"]).split("E100-")
|
|
if len(xpath) == 2 and os.path.isfile(xpath[0] + xpath[1]):
|
|
xpath = xpath[0] + xpath[1]
|
|
elif "fbv2" in str(data["last_checkpoint"]):
|
|
xpath = str(data["last_checkpoint"]).replace("fbv2", "mask_gumbel")
|
|
elif "tunas" in str(data["last_checkpoint"]):
|
|
xpath = str(data["last_checkpoint"]).replace("tunas", "mask_rl")
|
|
else:
|
|
raise ValueError(
|
|
"Invalid path: {:}".format(data["last_checkpoint"])
|
|
)
|
|
data = torch.load(xpath, 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
|
|
|
|
|
|
y_min_s = {
|
|
("cifar10", "tss"): 90,
|
|
("cifar10", "sss"): 92,
|
|
("cifar100", "tss"): 65,
|
|
("cifar100", "sss"): 65,
|
|
("ImageNet16-120", "tss"): 36,
|
|
("ImageNet16-120", "sss"): 40,
|
|
}
|
|
|
|
y_max_s = {
|
|
("cifar10", "tss"): 94.5,
|
|
("cifar10", "sss"): 93.3,
|
|
("cifar100", "tss"): 72,
|
|
("cifar100", "sss"): 70,
|
|
("ImageNet16-120", "tss"): 44,
|
|
("ImageNet16-120", "sss"): 46,
|
|
}
|
|
|
|
name2label = {
|
|
"cifar10": "CIFAR-10",
|
|
"cifar100": "CIFAR-100",
|
|
"ImageNet16-120": "ImageNet-16-120",
|
|
}
|
|
|
|
name2suffix = {
|
|
("sss", "warm"): "-WARM0.3",
|
|
("sss", "none"): "-WARMNone",
|
|
("tss", "none"): None,
|
|
("tss", None): None,
|
|
}
|
|
|
|
|
|
def visualize_curve(api, vis_save_dir, search_space, suffix):
|
|
vis_save_dir = vis_save_dir.resolve()
|
|
vis_save_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
dpi, width, height = 250, 5200, 1400
|
|
figsize = width / float(dpi), height / float(dpi)
|
|
LabelSize, LegendFontsize = 16, 16
|
|
|
|
def sub_plot_fn(ax, dataset):
|
|
print("{:} plot {:10s}".format(time_string(), dataset))
|
|
alg2data = fetch_data(
|
|
search_space=search_space,
|
|
dataset=dataset,
|
|
suffix=name2suffix[(search_space, suffix)],
|
|
)
|
|
alg2accuracies = OrderedDict()
|
|
epochs = 100
|
|
colors = ["b", "g", "c", "m", "y", "r"]
|
|
ax.set_xlim(0, epochs)
|
|
# ax.set_ylim(y_min_s[(dataset, search_space)], y_max_s[(dataset, search_space)])
|
|
for idx, (alg, data) in enumerate(alg2data.items()):
|
|
xs, accuracies = [], []
|
|
for iepoch in range(epochs + 1):
|
|
try:
|
|
structures, accs = [_[iepoch - 1] for _ in data], []
|
|
except:
|
|
raise ValueError(
|
|
"This alg {:} on {:} has invalid checkpoints.".format(
|
|
alg, dataset
|
|
)
|
|
)
|
|
for structure in structures:
|
|
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)
|
|
alg2accuracies[alg] = accuracies
|
|
ax.plot(xs, accuracies, c=colors[idx], label="{:}".format(alg))
|
|
ax.set_xlabel("The searching epoch", fontsize=LabelSize)
|
|
ax.set_ylabel(
|
|
"Test accuracy on {:}".format(name2label[dataset]), fontsize=LabelSize
|
|
)
|
|
ax.set_title(
|
|
"Searching results on {:}".format(name2label[dataset]),
|
|
fontsize=LabelSize + 4,
|
|
)
|
|
structures, valid_accs, test_accs = [_[epochs - 1] for _ in data], [], []
|
|
print(
|
|
"{:} plot alg : {:} -- final {:} architectures.".format(
|
|
time_string(), alg, len(structures)
|
|
)
|
|
)
|
|
for arch in structures:
|
|
valid_acc, test_acc, _ = get_valid_test_acc(api, arch, dataset)
|
|
test_accs.append(test_acc)
|
|
valid_accs.append(valid_acc)
|
|
print(
|
|
"{:} plot alg : {:} -- validation: {:.2f}$\pm${:.2f} -- test: {:.2f}$\pm${:.2f}".format(
|
|
time_string(),
|
|
alg,
|
|
np.mean(valid_accs),
|
|
np.std(valid_accs),
|
|
np.mean(test_accs),
|
|
np.std(test_accs),
|
|
)
|
|
)
|
|
ax.legend(loc=4, fontsize=LegendFontsize)
|
|
|
|
fig, axs = plt.subplots(1, 3, figsize=figsize)
|
|
datasets = ["cifar10", "cifar100", "ImageNet16-120"]
|
|
for dataset, ax in zip(datasets, axs):
|
|
sub_plot_fn(ax, dataset)
|
|
print("sub-plot {:} on {:} done.".format(dataset, search_space))
|
|
save_path = (
|
|
vis_save_dir / "{:}-ws-{:}-curve.png".format(search_space, suffix)
|
|
).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", formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
)
|
|
parser.add_argument(
|
|
"--save_dir",
|
|
type=str,
|
|
default="output/vis-nas-bench/nas-algos",
|
|
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)
|
|
visualize_curve(api_tss, save_dir, "tss", None)
|
|
|
|
api_sss = create(None, "sss", fast_mode=True, verbose=False)
|
|
visualize_curve(api_sss, save_dir, "sss", "warm")
|
|
visualize_curve(api_sss, save_dir, "sss", "none")
|