imagenet_mobilenet_search.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "intensive-warner",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys;\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "sys.path.append(os.path.join(os.getcwd(),'./'))\n",
    "\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xgboost as xgb\n",
    "from copy import deepcopy\n",
    "from tqdm import tqdm\n",
    "\n",
    "from ofa.utils import list_mean\n",
    "from ofa.imagenet_classification.data_providers.imagenet import ImagenetDataProvider\n",
    "from ofa.imagenet_classification.run_manager import ImagenetRunConfig, RunManager\n",
    "from ofa.model_zoo import ofa_net\n",
    "from ofa.nas.efficiency_predictor import MBv3LatencyTable\n",
    "from ofa.nas.accuracy_predictor.arch_encoder import MobileNetArchEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pleased-senegal",
   "metadata": {},
   "outputs": [],
   "source": [
    "__all__ = ['net_setting2id', 'net_id2setting', 'AccuracyDataset']\n",
    "\n",
    "def net_setting2id(net_setting):\n",
    "  return json.dumps(net_setting)\n",
    "\n",
    "def net_id2setting(net_id):\n",
    "  return json.loads(net_id)\n",
    "\n",
    "class AccuracyDataset:\n",
    "  def __init__(self, flops_constraint):\n",
    "    self.image_size = 236\n",
    "    batch_size = 100 # Adjust Test Batch Size here\n",
    "    ImagenetDataProvider.DEFAULT_PATH = '/ssd1/DATASET/nasbench201/imagenet' # Set Full ImageNet Path here\n",
    "    self.ofa_network = ofa_net('ofa_mbv3_d234_e346_k357_w1.2', pretrained=True)\n",
    "    self.run_config = ImagenetRunConfig(train_batch_size=batch_size,\n",
    "              test_batch_size=batch_size, \n",
    "              n_worker=8,\n",
    "              image_size=self.image_size)\n",
    "    self.run_manager = RunManager('.tmp/', self.ofa_network, self.run_config, init=False)\n",
    "    self.run_manager.run_config.data_provider.assign_active_img_size(self.image_size)\n",
    "    self.flops_constraint = flops_constraint\n",
    "\n",
    "  def random_sample(self, num, net_setting_list=None):\n",
    "    if net_setting_list is None:\n",
    "      net_id_list = set()\n",
    "      self.flops_dict = dict()\n",
    "      while len(net_id_list) < num:\n",
    "        net_setting = self.ofa_network.sample_active_subnet()\n",
    "        net_config = self.ofa_network.get_active_net_config()\n",
    "        flops = int(MBv3LatencyTable.count_flops_given_config(net_config, image_size=self.image_size))\n",
    "        net_setting_full = {**net_setting, 'r': self.image_size, 'flops': flops}\n",
    "        if flops >= self.flops_constraint[0] and flops <= self.flops_constraint[1]:\n",
    "          net_id = net_setting2id(net_setting_full)\n",
    "          net_id_list.add(net_id)\n",
    "          self.flops_dict[net_id] = flops\n",
    "      net_id_list = list(net_id_list)\n",
    "      net_id_list.sort()\n",
    "      return net_id_list\n",
    "    else:\n",
    "      net_id_list = set()\n",
    "      self.flops_dict = dict()\n",
    "      for net_setting in net_setting_list:\n",
    "        net_id = net_setting2id(net_setting)\n",
    "        self.ofa_network.set_active_subnet(**net_setting)\n",
    "        net_config = self.ofa_network.get_active_net_config()\n",
    "        flops = int(MBv3LatencyTable.count_flops_given_config(net_config, image_size=self.image_size))\n",
    "        if flops >= self.flops_constraint[0] and flops <= self.flops_constraint[1]:\n",
    "          self.flops_dict[net_id] = flops\n",
    "          net_id_list.add(net_id)\n",
    "      net_id_list = list(net_id_list)\n",
    "      net_id_list.sort()\n",
    "      return net_id_list\n",
    "\n",
    "  def get_acc(self, net_id):\n",
    "#     flops = self.flops_dict[net_id]\n",
    "#     net_setting = net_id2setting(net_id)\n",
    "#     key = net_setting2id({**net_setting, 'flops': flops})\n",
    "#     key = net_id\n",
    "    net_setting = net_id2setting(net_id)\n",
    "    net_setting_str = ','.join(['%s_%s' % (key, '%.1f' % list_mean(val) if isinstance(val, list) else val) for key, val in net_setting.items()])\n",
    "    acc_dict = dict()\n",
    "    self.ofa_network.set_active_subnet(**net_setting)\n",
    "    self.run_manager.reset_running_statistics(self.ofa_network)\n",
    "    loss, (top1, top5) = self.run_manager.validate(is_test=True, run_str=net_setting_str, net=self.ofa_network, data_loader=None, no_logs=True)\n",
    "    metric = (top1, top5)\n",
    "    flops = net_id2setting(net_id)['flops']\n",
    "    print(f'net: {net_setting_str}, r:{self.image_size}, flops:{flops}, top1:{metric[0]}, top5:{metric[1]}')\n",
    "    acc_dict.update({\n",
    "     net_id: metric\n",
    "    })\n",
    "    return metric[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sublime-ireland",
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_dataset = AccuracyDataset(flops_constraint=[550, 600])\n",
    "net_id_list = accuracy_dataset.random_sample(num=10000)\n",
    "acc_array = []\n",
    "arch_array = []\n",
    "flops_array = []\n",
    "arch_raw = []\n",
    "for net_id in net_id_list:\n",
    "  arch_raw.append(net_id)\n",
    "  key_dict = json.loads(net_id)\n",
    "#   acc = accuracy_dataset.get_acc(net_id)\n",
    "  arch = key_dict['ks']+key_dict['e']+key_dict['d']\n",
    "  acc = -1\n",
    "  flops = key_dict['flops']\n",
    "  arch_array.append(arch)\n",
    "  acc_array.append(acc)\n",
    "  flops_array.append(flops)\n",
    "  \n",
    "arch_array = np.array(arch_array)\n",
    "acc_array = np.array(acc_array)\n",
    "flops_array = np.array(flops_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "third-basic",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_traj(accuracy_dataset, ARCH_RAW, ARCH, LABEL, seed,\n",
    "             num_sample_train_list, num_sample_method,\n",
    "             top_list_acc, keep_old='all', n_trees=1000, max_depth=20):\n",
    "  \n",
    "  num_iteration = len(num_sample_train_list)\n",
    "  arch_dict_data = {}\n",
    "  if num_sample_method == 'uniform':\n",
    "    wgt_list = [1 / num_iteration for i in range(num_iteration)]\n",
    "  else:\n",
    "    assert ValueError\n",
    "  wgt_list_sum = sum(wgt_list)\n",
    "  wgt_list = [w / float(wgt_list_sum) for w in wgt_list]\n",
    "  all_index = np.arange(len(LABEL))\n",
    "  all_index_selected = deepcopy(all_index)\n",
    "  keep_index_train = np.array([]).astype(np.int)\n",
    "  df_dict_all = pd.DataFrame()\n",
    "  for i in range(arch_array.shape[-1]):\n",
    "    df_dict_all[f'arch_{i + 1}'] = [float(data[i]) for data in ARCH]\n",
    "  sample_list = []\n",
    "  sample_space_list = []\n",
    "  sample_space_exlude_list = []\n",
    "  predict_best_list = []\n",
    "  acc_pred_list = []\n",
    "  acc_gt_list = []\n",
    "  acc_best_list = []\n",
    "  train_acc_dict = {}\n",
    "  for z, (num_sample_train, top_acc) in enumerate(zip(num_sample_train_list, top_list_acc)):\n",
    "    acc_gt_list.append(LABEL)\n",
    "    if len(all_index_selected) == 0:\n",
    "      break  \n",
    "    if z == 0:\n",
    "      all_index_sample = all_index_selected\n",
    "      np.random.seed(seed+z)\n",
    "      train_index = np.random.choice(all_index_sample, size=min(num_sample_train, len(all_index_sample)), replace=False)\n",
    "      sample_space_list.append(all_index_selected)\n",
    "      sample_space_exlude_list.append(all_index_selected)\n",
    "      sample_list.append(train_index)\n",
    "    else:\n",
    "      sample_space_list.append(all_index_selected)\n",
    "      if keep_index_train.size != 0:\n",
    "        all_index_sample = all_index_selected[~np.isin(all_index_selected, keep_index_train)]\n",
    "      else:\n",
    "        all_index_sample = all_index_selected\n",
    "#       print(f'{z} sample space: {len(all_index_selected)}, sample space actual: {len(all_index_sample)}')\n",
    "      if num_sample_train == 0:\n",
    "        break\n",
    "      np.random.seed(seed+z+100)\n",
    "      train_index_sample = np.random.choice(all_index_sample, size=min(num_sample_train, len(all_index_sample)),\n",
    "                                            replace=False)\n",
    "      if keep_index_train.size != 0:\n",
    "        train_index = np.concatenate((keep_index_train, train_index_sample))\n",
    "      else:\n",
    "        train_index = train_index_sample\n",
    "        \n",
    "      sample_space_exlude_list.append(all_index_sample)\n",
    "      sample_list.append(train_index_sample)\n",
    "    \n",
    "    assert len(train_index) != 0\n",
    "    test_index = np.setdiff1d(all_index, train_index)\n",
    "    assert len(test_index) + len(train_index) == len(all_index)\n",
    "    assert len(test_index) != 0\n",
    "    params = {'booster': 'gbtree',\n",
    "              'max_depth': max_depth,\n",
    "              'objective': 'reg:squarederror'}\n",
    "    label_list = []\n",
    "    for index in train_index:\n",
    "      if index in train_acc_dict.keys():\n",
    "        label_list.append(train_acc_dict[index])\n",
    "      else:\n",
    "        net_id = ARCH_RAW[index]\n",
    "        acc = accuracy_dataset.get_acc(net_id)\n",
    "        label_list.append(acc)\n",
    "        train_acc_dict[index] = acc\n",
    "    dtrain = xgb.DMatrix(data=df_dict_all.iloc[train_index], label=label_list)\n",
    "    dall = xgb.DMatrix(data=df_dict_all, label=LABEL)\n",
    "    bst = xgb.train(params=params, dtrain=dtrain, num_boost_round=n_trees)\n",
    "    pred_all = bst.predict(dall)\n",
    "    index_optimal_pred = np.argmax(pred_all)\n",
    "    acc_pred_list.append(deepcopy(pred_all))\n",
    "    predict_best_list.append(index_optimal_pred)\n",
    "    all_index_by_acc = np.array([i for (i, _) in sorted(zip(np.arange(len(pred_all)), pred_all), key=lambda pair: pair[-1], reverse=True)]).astype(np.int)\n",
    "    dist_to_label = np.abs(pred_all-LABEL)\n",
    "    all_index_selected = all_index_by_acc[:top_acc]\n",
    "    if keep_old == 'none':\n",
    "      keep_index_train = np.array([]).astype(np.int)\n",
    "    elif keep_old == 'top':\n",
    "      keep_index_train = np.array([i for i in all_index_selected if i in train_index]).astype(np.int)\n",
    "    elif keep_old == 'all':\n",
    "      keep_index_train = train_index.astype(np.int)\n",
    "    del bst\n",
    "    acc_pred_best = LABEL[index_optimal_pred]\n",
    "    acc_best_list.append(acc_pred_best)\n",
    "  if index_optimal_pred in train_acc_dict.keys():\n",
    "    best_acc = train_acc_dict[index]\n",
    "  else:\n",
    "    net_id = ARCH_RAW[index_optimal_pred]\n",
    "    best_acc = accuracy_dataset.get_acc(net_id)\n",
    "  return best_acc, ARCH_RAW[index_optimal_pred]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eastern-chain",
   "metadata": {},
   "outputs": [],
   "source": [
    "ARCH = arch_array\n",
    "LABEL = acc_array\n",
    "ARCH_RAW = arch_raw\n",
    "num_sample_each_iterations = 100 # Set Number of Samples each Iterations\n",
    "num_iterations = 10 # Set Number of Iterations, 800 Queries: Set to 8, 1000 Queries: Set to 10\n",
    "repeat = 5\n",
    "top = 1000\n",
    "num_sample_train_list = [num_sample_each_iterations]*num_iterations\n",
    "top_list_acc = [top]*num_iterations\n",
    "seed = 0\n",
    "acc_list = []\n",
    "best_net_id_list = []\n",
    "for seed in tqdm(range(repeat)):\n",
    "  acc, best_net_id = run_traj(accuracy_dataset, ARCH_RAW, ARCH, LABEL, seed,\n",
    "           num_sample_train_list=num_sample_train_list, num_sample_method='uniform',\n",
    "           top_list_acc=top_list_acc, keep_old='all', n_trees=1000, max_depth=20)\n",
    "  acc_list.append(acc)\n",
    "  best_net_id_list.append(best_net_id)\n",
    "  net_setting = net_id2setting(best_net_id)\n",
    "#     net_setting_str = ','.join(['%s_%s' % (key, '%.1f' % list_mean(val) if isinstance(val, list) else val) for key, val in net_setting.items()])\n",
    "  print(f'Queries: {np.sum(num_sample_train_list)}, Best Arch: {net_setting}, Best SuperNet Acc: {acc}')\n",
    "acc_mean = np.mean(acc_list)\n",
    "acc_std = np.std(acc_list)\n",
    "print(f'Top {num_sample_each_iterations}/{top}, SuperNet Acc: {acc_mean:0.4f} +_({acc_std:0.4f})')"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}