{
 "nbformat": 4,
 "nbformat_minor": 2,
 "metadata": {
  "language_info": {
   "name": "python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "version": "3.6.10-final"
  },
  "orig_nbformat": 2,
  "file_extension": ".py",
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3,
  "kernelspec": {
   "name": "python361064bitnnilatestcondabff8d66a619a4d26af34fe0fe687c7b0",
   "display_name": "Python 3.6.10 64-bit ('nnilatest': conda)"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example Usages of NAS Benchmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint\n",
    "import time\n",
    "\n",
    "from nni.nas.benchmarks.nasbench101 import query_nb101_trial_stats\n",
    "from nni.nas.benchmarks.nasbench201 import query_nb201_trial_stats\n",
    "from nni.nas.benchmarks.nds import query_nds_trial_stats\n",
    "\n",
    "ti = time.time()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NAS-Bench-101"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "{'config': {'arch': {'input1': [0],\n                     'input2': [1],\n                     'input3': [2],\n                     'input4': [0],\n                     'input5': [0, 3, 4],\n                     'input6': [2, 5],\n                     'op1': 'conv3x3-bn-relu',\n                     'op2': 'maxpool3x3',\n                     'op3': 'conv3x3-bn-relu',\n                     'op4': 'conv3x3-bn-relu',\n                     'op5': 'conv1x1-bn-relu'},\n            'hash': '00005c142e6f48ac74fdcf73e3439874',\n            'id': 4,\n            'num_epochs': 108,\n            'num_vertices': 7},\n 'id': 10,\n 'parameters': 8.55553,\n 'test_acc': 92.11738705635071,\n 'train_acc': 100.0,\n 'training_time': 106147.67578125,\n 'valid_acc': 92.41786599159241}\n{'config': {'arch': {'input1': [0],\n                     'input2': [1],\n                     'input3': [2],\n                     'input4': [0],\n                     'input5': [0, 3, 4],\n                     'input6': [2, 5],\n                     'op1': 'conv3x3-bn-relu',\n                     'op2': 'maxpool3x3',\n                     'op3': 'conv3x3-bn-relu',\n                     'op4': 'conv3x3-bn-relu',\n                     'op5': 'conv1x1-bn-relu'},\n            'hash': '00005c142e6f48ac74fdcf73e3439874',\n            'id': 4,\n            'num_epochs': 108,\n            'num_vertices': 7},\n 'id': 11,\n 'parameters': 8.55553,\n 'test_acc': 91.90705418586731,\n 'train_acc': 100.0,\n 'training_time': 106095.05859375,\n 'valid_acc': 92.45793223381042}\n{'config': {'arch': {'input1': [0],\n                     'input2': [1],\n                     'input3': [2],\n                     'input4': [0],\n                     'input5': [0, 3, 4],\n                     'input6': [2, 5],\n                     'op1': 'conv3x3-bn-relu',\n                     'op2': 'maxpool3x3',\n                     'op3': 'conv3x3-bn-relu',\n                     'op4': 'conv3x3-bn-relu',\n                     'op5': 'conv1x1-bn-relu'},\n            'hash': '00005c142e6f48ac74fdcf73e3439874',\n            'id': 4,\n            'num_epochs': 108,\n            'num_vertices': 7},\n 'id': 12,\n 'parameters': 8.55553,\n 'test_acc': 92.15745329856873,\n 'train_acc': 100.0,\n 'training_time': 106138.55712890625,\n 'valid_acc': 93.04887652397156}\n"
    }
   ],
   "source": [
    "arch = {\n",
    "    'op1': 'conv3x3-bn-relu',\n",
    "    'op2': 'maxpool3x3',\n",
    "    'op3': 'conv3x3-bn-relu',\n",
    "    'op4': 'conv3x3-bn-relu',\n",
    "    'op5': 'conv1x1-bn-relu',\n",
    "    'input1': [0],\n",
    "    'input2': [1],\n",
    "    'input3': [2],\n",
    "    'input4': [0],\n",
    "    'input5': [0, 3, 4],\n",
    "    'input6': [2, 5]\n",
    "}\n",
    "for t in query_nb101_trial_stats(arch, 108):\n",
    "    pprint.pprint(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NAS-Bench-201"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "{'config': {'arch': {'0_1': 'avg_pool_3x3',\n                     '0_2': 'conv_1x1',\n                     '0_3': 'conv_1x1',\n                     '1_2': 'skip_connect',\n                     '1_3': 'skip_connect',\n                     '2_3': 'skip_connect'},\n            'dataset': 'cifar100',\n            'id': 7,\n            'num_cells': 5,\n            'num_channels': 16,\n            'num_epochs': 200},\n 'flops': 15.65322,\n 'id': 3,\n 'latency': 0.013182918230692545,\n 'ori_test_acc': 53.11,\n 'ori_test_evaluation_time': 1.0195916947864352,\n 'ori_test_loss': 1.7307863704681397,\n 'parameters': 0.135156,\n 'seed': 999,\n 'test_acc': 53.07999995727539,\n 'test_evaluation_time': 0.5097958473932176,\n 'test_loss': 1.731276072692871,\n 'train_acc': 57.82,\n 'train_loss': 1.5116578379058838,\n 'training_time': 2888.4371995925903,\n 'valid_acc': 53.14000000610351,\n 'valid_evaluation_time': 0.5097958473932176,\n 'valid_loss': 1.7302966793060304}\n{'config': {'arch': {'0_1': 'avg_pool_3x3',\n                     '0_2': 'conv_1x1',\n                     '0_3': 'conv_1x1',\n                     '1_2': 'skip_connect',\n                     '1_3': 'skip_connect',\n                     '2_3': 'skip_connect'},\n            'dataset': 'cifar100',\n            'id': 7,\n            'num_cells': 5,\n            'num_channels': 16,\n            'num_epochs': 200},\n 'flops': 15.65322,\n 'id': 7,\n 'latency': 0.013182918230692545,\n 'ori_test_acc': 51.93,\n 'ori_test_evaluation_time': 1.0195916947864352,\n 'ori_test_loss': 1.7572312774658203,\n 'parameters': 0.135156,\n 'seed': 777,\n 'test_acc': 51.979999938964845,\n 'test_evaluation_time': 0.5097958473932176,\n 'test_loss': 1.7429540189743042,\n 'train_acc': 57.578,\n 'train_loss': 1.5114233912658692,\n 'training_time': 2888.4371995925903,\n 'valid_acc': 51.88,\n 'valid_evaluation_time': 0.5097958473932176,\n 'valid_loss': 1.7715086591720581}\n{'config': {'arch': {'0_1': 'avg_pool_3x3',\n                     '0_2': 'conv_1x1',\n                     '0_3': 'conv_1x1',\n                     '1_2': 'skip_connect',\n                     '1_3': 'skip_connect',\n                     '2_3': 'skip_connect'},\n            'dataset': 'cifar100',\n            'id': 7,\n            'num_cells': 5,\n            'num_channels': 16,\n            'num_epochs': 200},\n 'flops': 15.65322,\n 'id': 11,\n 'latency': 0.013182918230692545,\n 'ori_test_acc': 53.38,\n 'ori_test_evaluation_time': 1.0195916947864352,\n 'ori_test_loss': 1.7281623031616211,\n 'parameters': 0.135156,\n 'seed': 888,\n 'test_acc': 53.67999998779297,\n 'test_evaluation_time': 0.5097958473932176,\n 'test_loss': 1.7327697801589965,\n 'train_acc': 57.792,\n 'train_loss': 1.5091403088760376,\n 'training_time': 2888.4371995925903,\n 'valid_acc': 53.08000000610352,\n 'valid_evaluation_time': 0.5097958473932176,\n 'valid_loss': 1.7235548280715942}\n"
    }
   ],
   "source": [
    "arch = {\n",
    "    '0_1': 'avg_pool_3x3',\n",
    "    '0_2': 'conv_1x1',\n",
    "    '1_2': 'skip_connect',\n",
    "    '0_3': 'conv_1x1',\n",
    "    '1_3': 'skip_connect',\n",
    "    '2_3': 'skip_connect'\n",
    "}\n",
    "for t in query_nb201_trial_stats(arch, 200, 'cifar100'):\n",
    "    pprint.pprint(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NDS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "{'best_test_acc': 90.48,\n 'best_train_acc': 96.356,\n 'best_train_loss': 0.116,\n 'config': {'base_lr': 0.1,\n            'cell_spec': {},\n            'dataset': 'cifar10',\n            'generator': 'random',\n            'id': 45505,\n            'model_family': 'residual_bottleneck',\n            'model_spec': {'bot_muls': [0.0, 0.25, 0.25, 0.25],\n                           'ds': [1, 16, 1, 4],\n                           'num_gs': [1, 2, 1, 2],\n                           'ss': [1, 1, 2, 2],\n                           'ws': [16, 64, 128, 16]},\n            'num_epochs': 100,\n            'proposer': 'resnext-a',\n            'weight_decay': 0.0005},\n 'final_test_acc': 90.39,\n 'final_train_acc': 96.298,\n 'final_train_loss': 0.116,\n 'flops': 69.890986,\n 'id': 45505,\n 'iter_time': 0.065,\n 'parameters': 0.083002,\n 'seed': 1}\n"
    }
   ],
   "source": [
    "model_spec = {\n",
    "    'bot_muls': [0.0, 0.25, 0.25, 0.25],\n",
    "    'ds': [1, 16, 1, 4],\n",
    "    'num_gs': [1, 2, 1, 2],\n",
    "    'ss': [1, 1, 2, 2],\n",
    "    'ws': [16, 64, 128, 16]\n",
    "}\n",
    "# Use none as a wildcard\n",
    "for t in query_nds_trial_stats('residual_bottleneck', None, None, model_spec, None, 'cifar10'):\n",
    "    pprint.pprint(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "{'best_test_acc': 93.58,\n 'best_train_acc': 99.772,\n 'best_train_loss': 0.011,\n 'config': {'base_lr': 0.1,\n            'cell_spec': {},\n            'dataset': 'cifar10',\n            'generator': 'random',\n            'id': 108998,\n            'model_family': 'residual_basic',\n            'model_spec': {'ds': [1, 12, 12, 12],\n                           'ss': [1, 1, 2, 2],\n                           'ws': [16, 24, 24, 40]},\n            'num_epochs': 100,\n            'proposer': 'resnet',\n            'weight_decay': 0.0005},\n 'final_test_acc': 93.49,\n 'final_train_acc': 99.772,\n 'final_train_loss': 0.011,\n 'flops': 184.519578,\n 'id': 108998,\n 'iter_time': 0.059,\n 'parameters': 0.594138,\n 'seed': 1}\n"
    }
   ],
   "source": [
    "model_spec = {'ds': [1, 12, 12, 12], 'ss': [1, 1, 2, 2], 'ws': [16, 24, 24, 40]}\n",
    "for t in query_nds_trial_stats('residual_basic', 'resnet', 'random', model_spec, {}, 'cifar10'):\n",
    "    pprint.pprint(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "{'best_test_acc': 84.5,\n 'best_train_acc': 89.66499999999999,\n 'best_train_loss': 0.302,\n 'config': {'base_lr': 0.1,\n            'cell_spec': {},\n            'dataset': 'cifar10',\n            'generator': 'random',\n            'id': 139492,\n            'model_family': 'vanilla',\n            'model_spec': {'ds': [1, 12, 12, 12],\n                           'ss': [1, 1, 2, 2],\n                           'ws': [16, 24, 32, 40]},\n            'num_epochs': 100,\n            'proposer': 'vanilla',\n            'weight_decay': 0.0005},\n 'final_test_acc': 84.35,\n 'final_train_acc': 89.633,\n 'final_train_loss': 0.303,\n 'flops': 208.36393,\n 'id': 154692,\n 'iter_time': 0.058,\n 'parameters': 0.68977,\n 'seed': 1}\n"
    }
   ],
   "source": [
    "# get the first one\n",
    "pprint.pprint(next(query_nds_trial_stats('vanilla', None, None, None, None, None)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "{'best_test_acc': 93.37,\n 'best_train_acc': 99.91,\n 'best_train_loss': 0.006,\n 'config': {'base_lr': 0.1,\n            'cell_spec': {'normal_0_input_x': 0,\n                          'normal_0_input_y': 1,\n                          'normal_0_op_x': 'avg_pool_3x3',\n                          'normal_0_op_y': 'conv_7x1_1x7',\n                          'normal_1_input_x': 2,\n                          'normal_1_input_y': 0,\n                          'normal_1_op_x': 'sep_conv_3x3',\n                          'normal_1_op_y': 'sep_conv_5x5',\n                          'normal_2_input_x': 2,\n                          'normal_2_input_y': 2,\n                          'normal_2_op_x': 'dil_sep_conv_3x3',\n                          'normal_2_op_y': 'dil_sep_conv_3x3',\n                          'normal_3_input_x': 4,\n                          'normal_3_input_y': 4,\n                          'normal_3_op_x': 'skip_connect',\n                          'normal_3_op_y': 'dil_sep_conv_3x3',\n                          'normal_4_input_x': 2,\n                          'normal_4_input_y': 4,\n                          'normal_4_op_x': 'conv_7x1_1x7',\n                          'normal_4_op_y': 'sep_conv_3x3',\n                          'normal_concat': [3, 5, 6],\n                          'reduce_0_input_x': 0,\n                          'reduce_0_input_y': 1,\n                          'reduce_0_op_x': 'avg_pool_3x3',\n                          'reduce_0_op_y': 'dil_sep_conv_3x3',\n                          'reduce_1_input_x': 0,\n                          'reduce_1_input_y': 0,\n                          'reduce_1_op_x': 'sep_conv_3x3',\n                          'reduce_1_op_y': 'sep_conv_3x3',\n                          'reduce_2_input_x': 2,\n                          'reduce_2_input_y': 0,\n                          'reduce_2_op_x': 'skip_connect',\n                          'reduce_2_op_y': 'sep_conv_7x7',\n                          'reduce_3_input_x': 4,\n                          'reduce_3_input_y': 4,\n                          'reduce_3_op_x': 'conv_7x1_1x7',\n                          'reduce_3_op_y': 'skip_connect',\n                          'reduce_4_input_x': 0,\n                          'reduce_4_input_y': 5,\n                          'reduce_4_op_x': 'conv_7x1_1x7',\n                          'reduce_4_op_y': 'conv_7x1_1x7',\n                          'reduce_concat': [3, 6]},\n            'dataset': 'cifar10',\n            'generator': 'random',\n            'id': 1,\n            'model_family': 'nas_cell',\n            'model_spec': {'aux': False,\n                           'depth': 12,\n                           'drop_prob': 0.0,\n                           'num_nodes_normal': 5,\n                           'num_nodes_reduce': 5,\n                           'width': 32},\n            'num_epochs': 100,\n            'proposer': 'amoeba',\n            'weight_decay': 0.0005},\n 'final_test_acc': 93.27,\n 'final_train_acc': 99.91,\n 'final_train_loss': 0.006,\n 'flops': 664.400586,\n 'id': 1,\n 'iter_time': 0.281,\n 'parameters': 4.190314,\n 'seed': 1}\n"
    }
   ],
   "source": [
    "# count number\n",
    "model_spec = {'num_nodes_normal': 5, 'num_nodes_reduce': 5, 'depth': 12, 'width': 32, 'aux': False, 'drop_prob': 0.0}\n",
    "cell_spec = {\n",
    "    'normal_0_op_x': 'avg_pool_3x3',\n",
    "    'normal_0_input_x': 0,\n",
    "    'normal_0_op_y': 'conv_7x1_1x7',\n",
    "    'normal_0_input_y': 1,\n",
    "    'normal_1_op_x': 'sep_conv_3x3',\n",
    "    'normal_1_input_x': 2,\n",
    "    'normal_1_op_y': 'sep_conv_5x5',\n",
    "    'normal_1_input_y': 0,\n",
    "    'normal_2_op_x': 'dil_sep_conv_3x3',\n",
    "    'normal_2_input_x': 2,\n",
    "    'normal_2_op_y': 'dil_sep_conv_3x3',\n",
    "    'normal_2_input_y': 2,\n",
    "    'normal_3_op_x': 'skip_connect',\n",
    "    'normal_3_input_x': 4,\n",
    "    'normal_3_op_y': 'dil_sep_conv_3x3',\n",
    "    'normal_3_input_y': 4,\n",
    "    'normal_4_op_x': 'conv_7x1_1x7',\n",
    "    'normal_4_input_x': 2,\n",
    "    'normal_4_op_y': 'sep_conv_3x3',\n",
    "    'normal_4_input_y': 4,\n",
    "    'normal_concat': [3, 5, 6],\n",
    "    'reduce_0_op_x': 'avg_pool_3x3',\n",
    "    'reduce_0_input_x': 0,\n",
    "    'reduce_0_op_y': 'dil_sep_conv_3x3',\n",
    "    'reduce_0_input_y': 1,\n",
    "    'reduce_1_op_x': 'sep_conv_3x3',\n",
    "    'reduce_1_input_x': 0,\n",
    "    'reduce_1_op_y': 'sep_conv_3x3',\n",
    "    'reduce_1_input_y': 0,\n",
    "    'reduce_2_op_x': 'skip_connect',\n",
    "    'reduce_2_input_x': 2,\n",
    "    'reduce_2_op_y': 'sep_conv_7x7',\n",
    "    'reduce_2_input_y': 0,\n",
    "    'reduce_3_op_x': 'conv_7x1_1x7',\n",
    "    'reduce_3_input_x': 4,\n",
    "    'reduce_3_op_y': 'skip_connect',\n",
    "    'reduce_3_input_y': 4,\n",
    "    'reduce_4_op_x': 'conv_7x1_1x7',\n",
    "    'reduce_4_input_x': 0,\n",
    "    'reduce_4_op_y': 'conv_7x1_1x7',\n",
    "    'reduce_4_input_y': 5,\n",
    "    'reduce_concat': [3, 6]\n",
    "}\n",
    "\n",
    "for t in query_nds_trial_stats('nas_cell', None, None, model_spec, cell_spec, 'cifar10'):\n",
    "    assert t['config']['model_spec'] == model_spec\n",
    "    assert t['config']['cell_spec'] == cell_spec\n",
    "    pprint.pprint(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "NDS (amoeba) count: 5107\n"
    }
   ],
   "source": [
    "# count number\n",
    "print('NDS (amoeba) count:', len(list(query_nds_trial_stats(None, 'amoeba', None, None, None, None, None))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "Elapsed time:  1.9107539653778076 seconds\n"
    }
   ],
   "source": [
    "print('Elapsed time: ', time.time() - ti, 'seconds')"
   ]
  }
 ]
}