{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retiarii Example - Multi-trial NAS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example will show Retiarii's ability to **express** and **explore** the model space for Neural Architecture Search and Hyper-Parameter Tuning in a simple way. The video demo is in [YouTube](https://youtu.be/eQUlABCO0o8) and [Bilibili](https://www.bilibili.com/video/BV14h411v7kZ/).\n",
    "\n",
    "Let's start the journey with Retiarii!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Express the Model Space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model space is defined by users to express a set of models that they want to explore, which contains potentially good-performing models. In Retiarii framework, a model space is defined with two parts: a base model and possible mutations on the base model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1.1: Define the Base Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining a base model is almost the same as defining a PyTorch (or TensorFlow) model. Usually, you only need to replace the code ``import torch.nn as nn`` with ``import nni.retiarii.nn.pytorch as nn`` to use NNI wrapped PyTorch modules. Below is a very simple example of defining a base model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import nni.retiarii.nn.pytorch as nn\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 3, padding=1)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 3, padding=1)\n",
    "        self.conv3 = nn.Conv2d(16, 16, 1)\n",
    "\n",
    "        self.bn = nn.BatchNorm2d(16)\n",
    "\n",
    "        self.gap = nn.AdaptiveAvgPool2d(4)\n",
    "        self.fc1 = nn.Linear(16 * 4 * 4, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        bs = x.size(0)\n",
    "\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x0 = F.relu(self.conv2(x))\n",
    "        x1 = F.relu(self.conv3(x0))\n",
    "\n",
    "        x1 += x0\n",
    "        x = self.pool(self.bn(x1))\n",
    "\n",
    "        x = self.gap(x).view(bs, -1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "    \n",
    "model = Net()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1.2: Define the Model Mutations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A base model is only one concrete model, not a model space. NNI provides APIs and primitives for users to express how the base model can be mutated, i.e., a model space that includes many models. The following will use inline Mutation APIs ``LayerChoice`` to choose a layer from candidate operations and use ``InputChoice`` to try out skip connection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import nni.retiarii.nn.pytorch as nn\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        # self.conv1 = nn.Conv2d(3, 6, 3, padding=1)\n",
    "        self.conv1 = nn.LayerChoice([nn.Conv2d(3, 6, 3, padding=1), nn.Conv2d(3, 6, 5, padding=2)])\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        # self.conv2 = nn.Conv2d(6, 16, 3, padding=1)\n",
    "        self.conv2 = nn.LayerChoice([nn.Conv2d(6, 16, 3, padding=1), nn.Conv2d(6, 16, 5, padding=2)])\n",
    "        self.conv3 = nn.Conv2d(16, 16, 1)\n",
    "\n",
    "        self.skipconnect = nn.InputChoice(n_candidates=2)\n",
    "        self.bn = nn.BatchNorm2d(16)\n",
    "\n",
    "        self.gap = nn.AdaptiveAvgPool2d(4)\n",
    "        self.fc1 = nn.Linear(16 * 4 * 4, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        bs = x.size(0)\n",
    "\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x0 = F.relu(self.conv2(x))\n",
    "        x1 = F.relu(self.conv3(x0))\n",
    "\n",
    "        x1 = self.skipconnect([x1, x1+x0])\n",
    "        x = self.pool(self.bn(x1))\n",
    "\n",
    "        x = self.gap(x).view(bs, -1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "    \n",
    "model = Net()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Explore the Model Space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will demo ths **multi-trial** NAS method first. In the multi-trial NAS process, the search strategy repeatedly generates new models, and the model evaluator is for training and validating each generated model. The obtained performance of a generated model is collected and sent to the search strategy for generating better models. \n",
    "\n",
    "Users can choose a proper search strategy to explore the model space, and use a chosen or user-defined model evaluator to evaluate the performance of each sampled model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.1: Choose or Write a Search Strategy\n",
    "\n",
    "Currently, Retiarii supports many common strategies, such as Random, Regularized evolution and TPE, etc. According to the APIs of Retiarii, you can customize a new strategy easily, and there we use the TPE strategy as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nni.retiarii.strategy as strategy\n",
    "\n",
    "simple_strategy = strategy.TPEStrategy() # choice: Random, GridSearch, RegularizedEvolution, TPEStrategy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.2: Choose or Write a Model Evaluator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model evaluator should correctly identify the use case of the model and the optimization goal. For example, on a classification task, an <input, label> dataset is needed, the loss function could be cross entropy and the optimized metric could be the accuracy.\n",
    "\n",
    "Retiarii provides two ways for users to write a new model evaluator. In the context of PyTorch, Retiarii has provided two built-in model evaluators, designed for simple use cases: classification and regression. These two evaluators are built upon the awesome library PyTorch-Lightning. Here we take a classification task on CIFAR10 dataset as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:15:27] INFO (lightning/MainThread) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:15:27] INFO (lightning/MainThread) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:15:27] INFO (lightning/MainThread) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    }
   ],
   "source": [
    "from torchvision import transforms\n",
    "from torchvision.datasets import CIFAR10\n",
    "from nni.retiarii import serialize\n",
    "import nni.retiarii.evaluator.pytorch.lightning as pl\n",
    "\n",
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
    "train_dataset = serialize(CIFAR10, root=\"./data\", train=True, download=True, transform=transform)\n",
    "test_dataset = serialize(CIFAR10, root=\"./data\", train=False, download=True, transform=transform)\n",
    "\n",
    "trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset, batch_size=64),\n",
    "                            val_dataloaders=pl.DataLoader(test_dataset, batch_size=64),\n",
    "                            max_epochs=2, gpus=[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.3: Configure the Experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After all the above are prepared, it is time to configure an experiment to do the model search. The basic experiment configuration is as follows, and advanced configuration reference on [this page](https://nni.readthedocs.io/en/stable/reference/experiment_config.html).\n",
    "\n",
    "NNI allows users to run experiments in different training platforms to speed up the search, like  Local Machine, Remote Servers, OpenPAI, Kubeflow, FrameworkController on K8S, DLWorkspace, Azure Machine Learning, AdaptDL, other cloud options, and even Hybrid mode. There we use the local mode with GPU speeding up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nni.retiarii.experiment.pytorch import RetiariiExeConfig, RetiariiExperiment\n",
    "\n",
    "exp = RetiariiExperiment(model, trainer, [], simple_strategy)\n",
    "\n",
    "exp_config = RetiariiExeConfig('local')\n",
    "exp_config.experiment_name = 'Retiarii example'\n",
    "exp_config.trial_concurrency = 2\n",
    "exp_config.max_trial_number = 10\n",
    "exp_config.trial_gpu_number = 2\n",
    "exp_config.max_experiment_duration = '5m'\n",
    "exp_config.training_service.use_active_gpu = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.4: Run and View the Experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can launch the experiment now! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:15:34] INFO (nni.experiment/MainThread) Creating experiment, Experiment ID: d9cseb3g\n",
      "[2021-06-07 11:15:34] INFO (nni.experiment/MainThread) Connecting IPC pipe...\n",
      "[2021-06-07 11:15:34] INFO (nni.experiment/MainThread) Statring web server...\n",
      "[2021-06-07 11:15:35] INFO (nni.experiment/MainThread) Setting up...\n",
      "[2021-06-07 11:15:36] INFO (nni.runtime.msg_dispatcher_base/Thread-6) Dispatcher started\n",
      "[2021-06-07 11:15:36] INFO (nni.retiarii.experiment.pytorch/MainThread) Web UI URLs: http://127.0.0.1:8745\n",
      "[2021-06-07 11:15:36] INFO (nni.retiarii.experiment.pytorch/MainThread) Start strategy...\n",
      "[2021-06-07 11:15:36] INFO (nni.retiarii.strategy.tpe_strategy/MainThread) TPE strategy has been started.\n",
      "[2021-06-07 11:15:36] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.001164 seconds\n",
      "[2021-06-07 11:15:36] INFO (hyperopt.tpe/MainThread) TPE using 0 trials\n",
      "[2021-06-07 11:15:36] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.001256 seconds\n",
      "[2021-06-07 11:15:36] INFO (hyperopt.tpe/MainThread) TPE using 0 trials\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:16:31] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:16:31] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:16:31] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:16:33] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002677 seconds\n",
      "[2021-06-07 11:16:33] INFO (hyperopt.tpe/MainThread) TPE using 1/1 trials with best loss 0.626600\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:16:36] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:16:36] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:16:36] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:16:37] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002730 seconds\n",
      "[2021-06-07 11:16:37] INFO (hyperopt.tpe/MainThread) TPE using 1/1 trials with best loss 0.626600\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:17:26] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:17:26] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:17:26] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:17:27] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.003051 seconds\n",
      "[2021-06-07 11:17:27] INFO (hyperopt.tpe/MainThread) TPE using 2/2 trials with best loss 0.594700\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:17:31] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:17:31] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:17:31] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:17:31] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002537 seconds\n",
      "[2021-06-07 11:17:31] INFO (hyperopt.tpe/MainThread) TPE using 3/3 trials with best loss 0.594700\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:18:21] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:18:21] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:18:21] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:18:22] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002532 seconds\n",
      "[2021-06-07 11:18:22] INFO (hyperopt.tpe/MainThread) TPE using 4/4 trials with best loss 0.594700\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:18:26] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:18:26] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:18:26] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:18:28] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002615 seconds\n",
      "[2021-06-07 11:18:28] INFO (hyperopt.tpe/MainThread) TPE using 6/6 trials with best loss 0.594700\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:19:16] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:19:16] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:19:16] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:19:18] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002395 seconds\n",
      "[2021-06-07 11:19:18] INFO (hyperopt.tpe/MainThread) TPE using 7/7 trials with best loss 0.594700\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:19:21] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:19:21] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:19:21] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:19:23] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002959 seconds\n",
      "[2021-06-07 11:19:23] INFO (hyperopt.tpe/MainThread) TPE using 7/7 trials with best loss 0.594700\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:20:12] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:20:12] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:20:12] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:20:13] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.003336 seconds\n",
      "[2021-06-07 11:20:13] INFO (hyperopt.tpe/MainThread) TPE using 8/8 trials with best loss 0.594700\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:20:22] INFO (lightning/Thread-5) GPU available: True, used: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:20:22] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-07 11:20:22] INFO (lightning/Thread-5) LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
      "[2021-06-07 11:20:22] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002093 seconds\n",
      "[2021-06-07 11:20:22] INFO (hyperopt.tpe/MainThread) TPE using 9/9 trials with best loss 0.593200\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "[2021-06-07 11:20:26] INFO (nni.retiarii.experiment.pytorch/Thread-7) Stopping experiment, please wait...\n",
      "[2021-06-07 11:20:26] INFO (nni.runtime.msg_dispatcher_base/Thread-6) Dispatcher exiting...\n",
      "[2021-06-07 11:20:26] INFO (nni.retiarii.experiment.pytorch/MainThread) Strategy exit\n",
      "[2021-06-07 11:20:26] INFO (nni.retiarii.experiment.pytorch/MainThread) Waiting for experiment to become DONE (you can ctrl+c if there is no running trial jobs)...\n",
      "[2021-06-07 11:20:27] INFO (nni.retiarii.experiment.pytorch/Thread-7) Experiment stopped\n",
      "[2021-06-07 11:20:29] INFO (nni.runtime.msg_dispatcher_base/Thread-6) Dispatcher terminiated\n"
     ]
    }
   ],
   "source": [
    "exp.run(exp_config, 8745)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides, NNI provides WebUI to help users view the experiment results and make more advanced analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.5: Export the top Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After searching, exporting the top model script is also very convenient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final model:\n",
      "import torch\n",
      "import torch.nn as nn\n",
      "import torch.nn.functional as F\n",
      "import torch.optim as optim\n",
      "\n",
      "import nni.retiarii.nn.pytorch\n",
      "\n",
      "import nni\n",
      "import torch\n",
      "\n",
      "\n",
      "class _model__conv1(nn.Module):\n",
      "    def __init__(self):\n",
      "        super().__init__()\n",
      "        self.layerchoice__mutation_1_0 = torch.nn.modules.conv.Conv2d(padding=1, in_channels=3, out_channels=6, kernel_size=3)\n",
      "\n",
      "    def forward(self, *_inputs):\n",
      "        layerchoice__mutation_1_0 = self.layerchoice__mutation_1_0(_inputs[0])\n",
      "        return layerchoice__mutation_1_0\n",
      "\n",
      "\n",
      "\n",
      "class _model__conv2(nn.Module):\n",
      "    def __init__(self):\n",
      "        super().__init__()\n",
      "        self.layerchoice__mutation_2_1 = torch.nn.modules.conv.Conv2d(padding=2, in_channels=6, out_channels=16, kernel_size=5)\n",
      "\n",
      "    def forward(self, *_inputs):\n",
      "        layerchoice__mutation_2_1 = self.layerchoice__mutation_2_1(_inputs[0])\n",
      "        return layerchoice__mutation_2_1\n",
      "\n",
      "\n",
      "\n",
      "class _model(nn.Module):\n",
      "    def __init__(self):\n",
      "        super().__init__()\n",
      "        self.__conv1 = _model__conv1()\n",
      "        self.__pool = torch.nn.modules.pooling.MaxPool2d(kernel_size=2, stride=2)\n",
      "        self.__conv2 = _model__conv2()\n",
      "        self.__conv3 = torch.nn.modules.conv.Conv2d(in_channels=16, out_channels=16, kernel_size=1)\n",
      "        self.__skipconnect = nni.retiarii.nn.pytorch.ChosenInputs(chosen=[1], reduction='sum')\n",
      "        self.__bn = torch.nn.modules.batchnorm.BatchNorm2d(num_features=16)\n",
      "        self.__gap = torch.nn.modules.pooling.AdaptiveAvgPool2d(output_size=4)\n",
      "        self.__fc1 = torch.nn.modules.linear.Linear(in_features=256, out_features=120)\n",
      "        self.__fc2 = torch.nn.modules.linear.Linear(in_features=120, out_features=84)\n",
      "        self.__fc3 = torch.nn.modules.linear.Linear(in_features=84, out_features=10)\n",
      "\n",
      "    def forward(self, x__1):\n",
      "        __Constant1 = -1\n",
      "        __Constant2 = 1\n",
      "        __Constant4 = False\n",
      "        __Constant5 = 0\n",
      "        __conv1 = self.__conv1(x__1)\n",
      "        __aten__size6 = x__1.size(dim=__Constant5)\n",
      "        __relu9 = F.relu(__conv1, __Constant4)\n",
      "        __ListConstruct21 = [__aten__size6, __Constant1]\n",
      "        __pool = self.__pool(__relu9)\n",
      "        __conv2 = self.__conv2(__pool)\n",
      "        __relu11 = F.relu(__conv2, __Constant4)\n",
      "        __conv3 = self.__conv3(__relu11)\n",
      "        __relu13 = F.relu(__conv3, __Constant4)\n",
      "        __aten__add15 = __relu13.add(other=__relu11, alpha=__Constant2)\n",
      "        __ListConstruct16 = [__relu13, __aten__add15]\n",
      "        __skipconnect = self.__skipconnect(__ListConstruct16)\n",
      "        __bn = self.__bn(__skipconnect)\n",
      "        __pool__19 = self.__pool(__bn)\n",
      "        __gap = self.__gap(__pool__19)\n",
      "        __aten__view22 = __gap.view(size=__ListConstruct21)\n",
      "        __fc1 = self.__fc1(__aten__view22)\n",
      "        __relu24 = F.relu(__fc1, __Constant4)\n",
      "        __fc2 = self.__fc2(__relu24)\n",
      "        __relu26 = F.relu(__fc2, __Constant4)\n",
      "        __fc3 = self.__fc3(__relu26)\n",
      "        return __fc3\n"
     ]
    }
   ],
   "source": [
    "print('Final model:')\n",
    "for model_code in exp.export_top_models():\n",
    "    print(model_code)"
   ]
  },
  {
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}