{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zero-Cost NAS Examples\n",
    "\n",
    "In this notebook, we provide examples of how to use zero-cost proxies within NAS algorithms. Specifically, we provide implementations of **random search** and **aging evolution search** with and without zero-cost warmup and move proposal.\n",
    "\n",
    "_note: the results in our ICLR paper were produced by a different (internal) AutoML tool that we are not publicly releasing. While it should be possible to get this notebook to exactly match our tool, we haven't done that here. We just provide examples to showcase the possible advantages of using zero-cost proxies._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "try to create the NAS-Bench-201 api from ../data/NAS-Bench-201-v1_0-e61699.pth\n"
     ]
    }
   ],
   "source": [
    "from nas_201_api import NASBench201API as API\n",
    "api = API('../data/NAS-Bench-201-v1_0-e61699.pth')\n",
    "api.verbose = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "synflow_proxy=[]\n",
    "f = open('../results_release/nasbench2/nb2_cf100_seed42_dlrandom_dlinfo1_initwnone_initbnone.p','rb')\n",
    "while(1):\n",
    "    try:\n",
    "        d = pickle.load(f)\n",
    "        synflow_proxy.append(d['logmeasures']['synflow'])\n",
    "    except EOFError:\n",
    "        break\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#dicts to map from index to spec or spec to index\n",
    "#spec is the model description within NAS, in our case we use an encoding [5, 5, 5, 5, 5, 5]\n",
    "#this is a length-6 vector, with each entry having a value between 0-4 \n",
    "#this is sufficient to describe any NAS-Bench-201 cell\n",
    "\n",
    "_opname_to_index = {\n",
    "    'none': 0,\n",
    "    'skip_connect': 1,\n",
    "    'nor_conv_1x1': 2,\n",
    "    'nor_conv_3x3': 3,\n",
    "    'avg_pool_3x3': 4\n",
    "}\n",
    "\n",
    "def get_spec_from_arch_str(arch_str):\n",
    "    nodes = arch_str.split('+')\n",
    "    nodes = [node[1:-1].split('|') for node in nodes]\n",
    "    nodes = [[op_and_input.split('~')[0]  for op_and_input in node] for node in nodes]\n",
    "\n",
    "    spec = [_opname_to_index[op] for node in nodes for op in node]\n",
    "    return spec\n",
    "\n",
    "idx_to_spec = {}\n",
    "for i, arch_str in enumerate(api):\n",
    "    idx_to_spec[i] = get_spec_from_arch_str(arch_str)\n",
    "\n",
    "spec_to_idx = {}\n",
    "for idx,spec in idx_to_spec.items():\n",
    "    spec_to_idx[str(spec)] = idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import copy\n",
    "random.seed = 42\n",
    "def random_spec():\n",
    "    return random.choice(list(idx_to_spec.values()))\n",
    "\n",
    "def mutate_spec(old_spec):\n",
    "    idx_to_change = random.randrange(len(old_spec))\n",
    "    entry_to_change = old_spec[idx_to_change]\n",
    "    possible_entries = [x for x in range(5) if x != entry_to_change]\n",
    "    new_entry = random.choice(possible_entries)\n",
    "    new_spec = copy.copy(old_spec)\n",
    "    new_spec[idx_to_change] = new_entry\n",
    "    return new_spec\n",
    "\n",
    "def mutate_spec_zero_cost(old_spec):\n",
    "    possible_specs = []\n",
    "    for idx_to_change in range(len(old_spec)): \n",
    "        entry_to_change = old_spec[idx_to_change]\n",
    "        possible_entries = [x for x in range(5) if x != entry_to_change]\n",
    "        for new_entry in possible_entries:\n",
    "            new_spec = copy.copy(old_spec)\n",
    "            new_spec[idx_to_change] = new_entry\n",
    "            possible_specs.append((synflow_proxy[spec_to_idx[str(new_spec)]], new_spec))\n",
    "    best_new_spec = sorted(possible_specs, key=lambda i:i[0])[-1][1]\n",
    "    if random.random() > 0.75:\n",
    "        best_new_spec = random.choice(possible_specs)[1]\n",
    "    return best_new_spec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_combination(iterable, sample_size):\n",
    "    pool = tuple(iterable)\n",
    "    n = len(pool)\n",
    "    indices = sorted(random.sample(range(n), sample_size))\n",
    "    return tuple(pool[i] for i in indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_evolution_search(max_trained_models=1000, \n",
    "                         pool_size=64, \n",
    "                         tournament_size=10, \n",
    "                         zero_cost_warmup=0, \n",
    "                         zero_cost_move=False):\n",
    "    \n",
    "    best_valids, best_tests = [0.0], [0.0]\n",
    "    pool = []   # (validation, spec) tuples\n",
    "    num_trained_models = 0\n",
    "\n",
    "    # fill the initial pool\n",
    "    if zero_cost_warmup > 0:\n",
    "        zero_cost_pool = []\n",
    "        for _ in range(zero_cost_warmup):\n",
    "            spec = random_spec()\n",
    "            spec_idx = spec_to_idx[str(spec)]\n",
    "            zero_cost_pool.append((synflow_proxy[spec_idx], spec))\n",
    "            zero_cost_pool = sorted(zero_cost_pool, key=lambda i:i[0], reverse=True)\n",
    "    for i in range(pool_size):\n",
    "        if zero_cost_warmup > 0:\n",
    "            spec = zero_cost_pool[i][1]\n",
    "        else:\n",
    "            spec = random_spec()\n",
    "        info = api.get_more_info(spec_to_idx[str(spec)], 'cifar100', iepoch=None, hp='200', is_random=False)\n",
    "        num_trained_models += 1\n",
    "        pool.append((info['valid-accuracy'], spec))\n",
    "\n",
    "        if info['valid-accuracy'] > best_valids[-1]:\n",
    "            best_valids.append(info['valid-accuracy'])\n",
    "        else:\n",
    "            best_valids.append(best_valids[-1])\n",
    "            \n",
    "        if info['test-accuracy'] > best_tests[-1]:\n",
    "            best_tests.append(info['test-accuracy'])\n",
    "        else:\n",
    "            best_tests.append(best_tests[-1])\n",
    "\n",
    "    # After the pool is seeded, proceed with evolving the population.\n",
    "    while(1):\n",
    "        sample = random_combination(pool, tournament_size)\n",
    "        best_spec = sorted(sample, key=lambda i:i[0])[-1][1]\n",
    "        if zero_cost_move:\n",
    "            new_spec = mutate_spec_zero_cost(best_spec)\n",
    "        else:\n",
    "            new_spec = mutate_spec(best_spec)\n",
    "\n",
    "        info = api.get_more_info(spec_to_idx[str(new_spec)], 'cifar100', iepoch=None, hp='200', is_random=False)\n",
    "        num_trained_models += 1\n",
    "\n",
    "        # kill the oldest individual in the population.\n",
    "        pool.append((info['valid-accuracy'], new_spec))\n",
    "        pool.pop(0)\n",
    "\n",
    "        if info['valid-accuracy'] > best_valids[-1]:\n",
    "            best_valids.append(info['valid-accuracy'])\n",
    "        else:\n",
    "            best_valids.append(best_valids[-1])\n",
    "            \n",
    "        if info['test-accuracy'] > best_tests[-1]:\n",
    "            best_tests.append(info['test-accuracy'])\n",
    "        else:\n",
    "            best_tests.append(best_tests[-1])\n",
    "\n",
    "        if num_trained_models >= max_trained_models:\n",
    "            break\n",
    "    best_tests.pop(0)\n",
    "    best_valids.pop(0)\n",
    "    return best_valids, best_tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_random_search(max_trained_models=1000, \n",
    "                      zero_cost_warmup=0):\n",
    "    \n",
    "    best_valids, best_tests = [0.0], [0.0]\n",
    "    pool = []   # (validation, spec) tuples\n",
    "    num_trained_models = 0\n",
    "\n",
    "    # fill the initial pool\n",
    "    if zero_cost_warmup > 0:\n",
    "        zero_cost_pool = []\n",
    "        for _ in range(zero_cost_warmup):\n",
    "            spec = random_spec()\n",
    "            spec_idx = spec_to_idx[str(spec)]\n",
    "            zero_cost_pool.append((synflow_proxy[spec_idx], spec))\n",
    "            zero_cost_pool = sorted(zero_cost_pool, key=lambda i:i[0], reverse=True)\n",
    "    for i in range(max_trained_models):\n",
    "        if i < zero_cost_warmup:\n",
    "            spec = zero_cost_pool[i][1]\n",
    "        else:\n",
    "            spec = random_spec()\n",
    "        info = api.get_more_info(spec_to_idx[str(spec)], 'cifar100', iepoch=None, hp='200', is_random=False)\n",
    "\n",
    "        if info['valid-accuracy'] > best_valids[-1]:\n",
    "            best_valids.append(info['valid-accuracy'])\n",
    "        else:\n",
    "            best_valids.append(best_valids[-1])\n",
    "            \n",
    "        if info['test-accuracy'] > best_tests[-1]:\n",
    "            best_tests.append(info['test-accuracy'])\n",
    "        else:\n",
    "            best_tests.append(best_tests[-1])\n",
    "            \n",
    "    best_tests.pop(0)\n",
    "    best_valids.pop(0)\n",
    "    return best_valids, best_tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:18<00:00,  1.80s/it]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "num_rounds = 10\n",
    "length = 300\n",
    "ae, ae_warmup, ae_move, rand, rand_warmup = [], [], [], [], []\n",
    "for _ in tqdm(range(num_rounds)):\n",
    "    ae_best_valids, ae_best_tests = run_evolution_search(max_trained_models=length)\n",
    "    ae.append(ae_best_tests)\n",
    "    ae_warmup_best_valids, ae_warmup_best_tests = run_evolution_search(max_trained_models=length, zero_cost_warmup=3000)\n",
    "    ae_warmup.append(ae_warmup_best_tests)\n",
    "    ae_move_best_valids, ae_move_best_tests = run_evolution_search(max_trained_models=length, zero_cost_move=True)\n",
    "    ae_move.append(ae_move_best_tests)\n",
    "    rand_best_valids, rand_best_tests = run_random_search(max_trained_models=length)\n",
    "    rand.append(rand_best_tests)\n",
    "    rand_warmup_best_valids, rand_warmup_best_tests = run_random_search(max_trained_models=length, zero_cost_warmup=3000)\n",
    "    rand_warmup.append(rand_warmup_best_tests)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def plot_experiment(exp_list, title):\n",
    "    def plot_exp(exp, label):\n",
    "        exp = np.array(exp) \n",
    "        q_75 = np.quantile(exp, .75, axis=0)\n",
    "        q_25 = np.quantile(exp, .25, axis=0)\n",
    "        mean = np.mean(exp, axis=0)\n",
    "        plt.plot(mean, label=label)\n",
    "        plt.fill_between(range(len(q_25)), q_25, q_75, alpha=0.1)\n",
    "    for exp,ename in exp_list:\n",
    "        plot_exp(exp,ename)\n",
    "    plt.grid()\n",
    "    plt.xlabel('Trained Models')\n",
    "    plt.ylabel('Test Accuracy')\n",
    "    plt.ylim(70,73.6)\n",
    "    plt.legend()\n",
    "    plt.title(title)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29584 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_experiment([(ae,'AE'), (ae_warmup,'AE + warmup (3000)'), (ae_move,'AE + move')], 'Aging Evolution Search')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21984 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_experiment([(rand,'RAND'), (rand_warmup,'RAND + warmup (3000)')], 'Random Search')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
