xautodl/exps/NATS-Bench/draw-ranks.py

###############################################################
# 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

from xautodl.config_utils import dict2config, load_config
from xautodl.log_utils import time_string
from xautodl.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__))