{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tabular Data Classification with NNI in AML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This simple example is to use NNI NAS 2.0(Retiarii) framework to search for the best neural architecture for tabular data classification task in Azure Machine Learning training platform.\n",
    "\n",
    "The video demo is in [YouTube](https://www.youtube.com/watch?v=PDVqBmm7Cro) and [Bilibili](https://www.bilibili.com/video/BV1oy4y1W7GF)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Prepare the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is to prepare the dataset. Here we use the Titanic dataset as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import TitanicDataset\n",
    "from nni.retiarii import serialize\n",
    "\n",
    "train_dataset = serialize(TitanicDataset, root='./data', train=True)\n",
    "test_dataset = serialize(TitanicDataset, root='./data', train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Define 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(NNI NAS 2.0) 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 2.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 nni.retiarii.nn.pytorch as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Net(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size):\n",
    "        super().__init__()\n",
    "\n",
    "        self.fc1 = nn.Linear(input_size, 16)\n",
    "        self.bn1 = nn.BatchNorm1d(16)\n",
    "        self.dropout1 = nn.Dropout(0.0)\n",
    "\n",
    "        self.fc2 = nn.Linear(16, 16)\n",
    "        self.bn2 = nn.BatchNorm1d(16)\n",
    "        self.dropout2 = nn.Dropout(0.0)\n",
    "\n",
    "        self.fc3 = nn.Linear(16, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.dropout1(F.relu(self.bn1(self.fc1(x))))\n",
    "        x = self.dropout2(F.relu(self.bn2(self.fc2(x))))\n",
    "        x = F.sigmoid(self.fc3(x))\n",
    "        return x\n",
    "    \n",
    "model_space = Net(len(train_dataset.__getitem__(0)[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.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 as a simple example. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nni.retiarii.nn.pytorch as nn\n",
    "import torch.nn.functional as F\n",
    "from nni.retiarii import model_wrapper\n",
    "\n",
    "@model_wrapper\n",
    "class Net(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size):\n",
    "        super().__init__()\n",
    "\n",
    "        self.hidden_dim1 = nn.ValueChoice(\n",
    "            [16, 32, 64, 128, 256, 512, 1024], label='hidden_dim1')\n",
    "        self.hidden_dim2 = nn.ValueChoice(\n",
    "            [16, 32, 64, 128, 256, 512, 1024], label='hidden_dim2')\n",
    "\n",
    "        self.fc1 = nn.Linear(input_size, self.hidden_dim1)\n",
    "        self.bn1 = nn.BatchNorm1d(self.hidden_dim1)\n",
    "        self.dropout1 = nn.Dropout(nn.ValueChoice([0.0, 0.25, 0.5]))\n",
    "\n",
    "        self.fc2 = nn.Linear(self.hidden_dim1, self.hidden_dim2)\n",
    "        self.bn2 = nn.BatchNorm1d(self.hidden_dim2)\n",
    "        self.dropout2 = nn.Dropout(nn.ValueChoice([0.0, 0.25, 0.5]))\n",
    "\n",
    "        self.fc3 = nn.Linear(self.hidden_dim2, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.dropout1(F.relu(self.bn1(self.fc1(x))))\n",
    "        x = self.dropout2(F.relu(self.bn2(self.fc2(x))))\n",
    "        x = F.sigmoid(self.fc3(x))\n",
    "        return x\n",
    "\n",
    "model_space = Net(len(train_dataset.__getitem__(0)[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides inline mutations, Retiarii also provides ``mutator``, a more general approach to express complex model space."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Explore the Defined Model Space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the 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 3.1: Choose a Search Strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nni.retiarii.strategy as strategy\n",
    "\n",
    "simple_strategy = strategy.TPEStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3.2: Choose or Write a Model Evaluator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 09:56:10] INFO (lightning/MainThread) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 09:56:10] INFO (lightning/MainThread) TPU available: None, using: 0 TPU cores\n"
     ]
    }
   ],
   "source": [
    "import nni.retiarii.evaluator.pytorch.lightning as pl\n",
    "\n",
    "trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset, batch_size=16),\n",
    "                                val_dataloaders=pl.DataLoader(\n",
    "                                test_dataset, batch_size=16),\n",
    "                                max_epochs=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: 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)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nni.retiarii.experiment.pytorch import RetiariiExeConfig, RetiariiExperiment\n",
    "\n",
    "exp = RetiariiExperiment(model_space, trainer, [], simple_strategy)\n",
    "\n",
    "exp_config = RetiariiExeConfig('aml')\n",
    "exp_config.experiment_name = 'titanic_example'\n",
    "exp_config.trial_concurrency = 2\n",
    "exp_config.max_trial_number = 20\n",
    "exp_config.max_experiment_duration = '2h'\n",
    "exp_config.nni_manager_ip = '' # your nni_manager_ip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before running experiments on AML(Azure Machine Learning) training service, you need to set up corresponding environment(refer to [AML mode doc](https://nni.readthedocs.io/en/stable/TrainingService/AMLMode.html)) and configure the following additional fields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Authenticate to your Azure subscription from the CLI.\n",
    "# If you have finished, please skip it.\n",
    "!az login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_config.training_service.subscription_id = '' # your subscription id\n",
    "exp_config.training_service.resource_group = '' # your resource group\n",
    "exp_config.training_service.workspace_name = '' # your workspace name\n",
    "exp_config.training_service.compute_target = '' # your compute target\n",
    "exp_config.training_service.docker_image = 'msranni/nni:latest'  # your docker image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Run and View the Experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can launch the experiment now! \n",
    "\n",
    "Besides, NNI provides WebUI to help users view the experiment results and make more advanced analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 09:56:54] INFO (nni.experiment/MainThread) Creating experiment, Experiment ID: 46den9qr\n",
      "[2021-06-08 09:56:55] INFO (nni.experiment/MainThread) Connecting IPC pipe...\n",
      "[2021-06-08 09:56:58] INFO (nni.experiment/MainThread) Starting web server...\n",
      "[2021-06-08 09:57:00] INFO (nni.experiment/MainThread) Setting up...\n",
      "[2021-06-08 09:57:05] INFO (nni.runtime.msg_dispatcher_base/Thread-8) Dispatcher started\n",
      "[2021-06-08 09:57:05] INFO (nni.retiarii.experiment.pytorch/MainThread) Web UI URLs: http://127.0.0.1:8745\n",
      "[2021-06-08 09:57:05] INFO (nni.retiarii.experiment.pytorch/MainThread) Start strategy...\n",
      "[2021-06-08 09:57:05] INFO (nni.retiarii.strategy.tpe_strategy/MainThread) TPE strategy has been started.\n",
      "[2021-06-08 09:57:05] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.001999 seconds\n",
      "[2021-06-08 09:57:05] INFO (hyperopt.tpe/MainThread) TPE using 0 trials\n",
      "[2021-06-08 09:57:10] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002029 seconds\n",
      "[2021-06-08 09:57:10] INFO (hyperopt.tpe/MainThread) TPE using 0 trials\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:03:55] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:03:55] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:03:56] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.000000 seconds\n",
      "[2021-06-08 10:03:56] INFO (hyperopt.tpe/MainThread) TPE using 1/1 trials with best loss 0.795455\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:04:46] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:04:46] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:04:46] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.000000 seconds\n",
      "[2021-06-08 10:04:46] INFO (hyperopt.tpe/MainThread) TPE using 2/2 trials with best loss 0.795455\n",
      "[2021-06-08 10:04:50] WARNING (nni.runtime.msg_dispatcher_base/Thread-8) assessor queue length: 20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:04:51] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:04:51] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:04:52] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.000000 seconds\n",
      "[2021-06-08 10:04:52] INFO (hyperopt.tpe/MainThread) TPE using 3/3 trials with best loss 0.795455\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:05:46] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:05:46] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:05:48] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002999 seconds\n",
      "[2021-06-08 10:05:48] INFO (hyperopt.tpe/MainThread) TPE using 4/4 trials with best loss 0.791667\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:05:56] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:05:56] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:05:56] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.000000 seconds\n",
      "[2021-06-08 10:05:56] INFO (hyperopt.tpe/MainThread) TPE using 5/5 trials with best loss 0.791667\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:06:26] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:06:26] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:06:27] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.004991 seconds\n",
      "[2021-06-08 10:06:27] INFO (hyperopt.tpe/MainThread) TPE using 6/6 trials with best loss 0.791667\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:07:06] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:07:06] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:07:07] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.006043 seconds\n",
      "[2021-06-08 10:07:07] INFO (hyperopt.tpe/MainThread) TPE using 7/7 trials with best loss 0.784091\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:07:56] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:07:56] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:07:57] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.006004 seconds\n",
      "[2021-06-08 10:07:57] INFO (hyperopt.tpe/MainThread) TPE using 8/8 trials with best loss 0.731061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:08:01] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:08:01] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:08:01] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.005000 seconds\n",
      "[2021-06-08 10:08:01] INFO (hyperopt.tpe/MainThread) TPE using 9/9 trials with best loss 0.731061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:08:56] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:08:56] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:08:58] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.004962 seconds\n",
      "[2021-06-08 10:08:58] INFO (hyperopt.tpe/MainThread) TPE using 10/10 trials with best loss 0.731061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:09:01] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:09:01] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:09:03] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.003043 seconds\n",
      "[2021-06-08 10:09:03] INFO (hyperopt.tpe/MainThread) TPE using 11/11 trials with best loss 0.731061\n",
      "[2021-06-08 10:10:24] WARNING (nni.runtime.msg_dispatcher_base/Thread-8) assessor queue length: 20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:10:27] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:10:27] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:10:28] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002005 seconds\n",
      "[2021-06-08 10:10:28] INFO (hyperopt.tpe/MainThread) TPE using 12/12 trials with best loss 0.731061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:10:52] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:10:52] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:10:53] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.122046 seconds\n",
      "[2021-06-08 10:10:53] INFO (hyperopt.tpe/MainThread) TPE using 13/13 trials with best loss 0.731061\n",
      "[2021-06-08 10:14:52] WARNING (nni.runtime.msg_dispatcher_base/Thread-8) assessor queue length: 20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:14:52] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:14:52] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:14:53] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002038 seconds\n",
      "[2021-06-08 10:14:53] INFO (hyperopt.tpe/MainThread) TPE using 14/14 trials with best loss 0.731061\n",
      "[2021-06-08 10:14:56] WARNING (nni.runtime.msg_dispatcher_base/Thread-8) assessor queue length: 20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:14:57] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:14:57] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:14:58] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.005870 seconds\n",
      "[2021-06-08 10:14:58] INFO (hyperopt.tpe/MainThread) TPE using 15/15 trials with best loss 0.731061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:16:07] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:16:07] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:16:08] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.004999 seconds\n",
      "[2021-06-08 10:16:08] INFO (hyperopt.tpe/MainThread) TPE using 16/16 trials with best loss 0.712121\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:16:48] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:16:48] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:16:48] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002000 seconds\n",
      "[2021-06-08 10:16:48] INFO (hyperopt.tpe/MainThread) TPE using 17/17 trials with best loss 0.712121\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:16:53] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:16:53] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:16:55] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002010 seconds\n",
      "[2021-06-08 10:16:55] INFO (hyperopt.tpe/MainThread) TPE using 18/18 trials with best loss 0.712121\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:17:43] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:17:43] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:17:44] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.006001 seconds\n",
      "[2021-06-08 10:17:44] INFO (hyperopt.tpe/MainThread) TPE using 19/19 trials with best loss 0.712121\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:18:03] INFO (lightning/Thread-5) GPU available: True, used: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TPU available: None, using: 0 TPU cores\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-08 10:18:03] INFO (lightning/Thread-5) TPU available: None, using: 0 TPU cores\n",
      "[2021-06-08 10:18:04] INFO (hyperopt.tpe/MainThread) tpe_transform took 0.002009 seconds\n",
      "[2021-06-08 10:18:04] INFO (hyperopt.tpe/MainThread) TPE using 20/20 trials with best loss 0.712121\n",
      "[2021-06-08 10:18:12] INFO (nni.retiarii.experiment.pytorch/Thread-9) Stopping experiment, please wait...\n",
      "[2021-06-08 10:18:14] INFO (nni.runtime.msg_dispatcher_base/Thread-8) Dispatcher exiting...\n",
      "[2021-06-08 10:18:14] INFO (nni.retiarii.experiment.pytorch/MainThread) Strategy exit\n",
      "[2021-06-08 10:18:14] 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-08 10:18:15] INFO (nni.retiarii.experiment.pytorch/Thread-9) Experiment stopped\n",
      "[2021-06-08 10:18:16] INFO (nni.runtime.msg_dispatcher_base/Thread-8) Dispatcher terminiated\n"
     ]
    }
   ],
   "source": [
    "exp.run(exp_config, 8745)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Export the top Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exporting the top model script is also very convenient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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 torch\n",
      "\n",
      "\n",
      "class _model(nn.Module):\n",
      "    def __init__(self):\n",
      "        super().__init__()\n",
      "        self.__fc1 = torch.nn.modules.linear.Linear(in_features=9, out_features=512)\n",
      "        self.__bn1 = torch.nn.modules.batchnorm.BatchNorm1d(num_features=512)\n",
      "        self.__dropout1 = torch.nn.modules.dropout.Dropout(p=0.0)\n",
      "        self.__fc2 = torch.nn.modules.linear.Linear(in_features=512, out_features=128)\n",
      "        self.__bn2 = torch.nn.modules.batchnorm.BatchNorm1d(num_features=128)\n",
      "        self.__dropout2 = torch.nn.modules.dropout.Dropout(p=0.25)\n",
      "        self.__fc3 = torch.nn.modules.linear.Linear(in_features=128, out_features=2)\n",
      "\n",
      "    def forward(self, x__1):\n",
      "        __Constant3 = False\n",
      "        __fc1 = self.__fc1(x__1)\n",
      "        __bn1 = self.__bn1(__fc1)\n",
      "        __relu7 = F.relu(__bn1, __Constant3)\n",
      "        __dropout1 = self.__dropout1(__relu7)\n",
      "        __fc2 = self.__fc2(__dropout1)\n",
      "        __bn2 = self.__bn2(__fc2)\n",
      "        __relu11 = F.relu(__bn2, __Constant3)\n",
      "        __dropout2 = self.__dropout2(__relu11)\n",
      "        __fc3 = self.__fc3(__dropout2)\n",
      "        __sigmoid13 = F.sigmoid(__fc3)\n",
      "        return __sigmoid13\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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}