{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Hello, NAS!\n\nThis is the 101 tutorial of Neural Architecture Search (NAS) on NNI.\nIn this tutorial, we will search for a neural architecture on MNIST dataset with the help of NAS framework of NNI, i.e., *Retiarii*.\nWe use multi-trial NAS as an example to show how to construct and explore a model space.\n\nThere are mainly three crucial components for a neural architecture search task, namely,\n\n* Model search space that defines a set of models to explore.\n* A proper strategy as the method to explore this model space.\n* A model evaluator that reports the performance of every model in the space.\n\nCurrently, PyTorch is the only supported framework by Retiarii, and we have only tested **PyTorch 1.7 to 1.10**.\nThis tutorial assumes PyTorch context but it should also apply to other frameworks, which is in our future plan.\n\n## Define your Model Space\n\nModel space is defined by users to express a set of models that users want to explore, which contains potentially good-performing models.\nIn this framework, a model space is defined with two parts: a base model and possible mutations on the base model.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Define Base Model\n\nDefining a base model is almost the same as defining a PyTorch (or TensorFlow) model.\nUsually, you only need to replace the code ``import torch.nn as nn`` with\n``import nni.retiarii.nn.pytorch as nn`` to use our wrapped PyTorch modules.\n\nBelow is a very simple example of defining a base model.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import torch\nimport torch.nn.functional as F\nimport nni.retiarii.nn.pytorch as nn\nfrom nni.retiarii import model_wrapper\n\n\n@model_wrapper # this decorator should be put on the out most\nclass Net(nn.Module):\n def __init__(self):\n super().__init__()\n self.conv1 = nn.Conv2d(1, 32, 3, 1)\n self.conv2 = nn.Conv2d(32, 64, 3, 1)\n self.dropout1 = nn.Dropout(0.25)\n self.dropout2 = nn.Dropout(0.5)\n self.fc1 = nn.Linear(9216, 128)\n self.fc2 = nn.Linear(128, 10)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n x = F.max_pool2d(self.conv2(x), 2)\n x = torch.flatten(self.dropout1(x), 1)\n x = self.fc2(self.dropout2(F.relu(self.fc1(x))))\n output = F.log_softmax(x, dim=1)\n return output" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ".. tip:: Always keep in mind that you should use ``import nni.retiarii.nn.pytorch as nn`` and :meth:`nni.retiarii.model_wrapper`.\n Many mistakes are a result of forgetting one of those.\n Also, please use ``torch.nn`` for submodules of ``nn.init``, e.g., ``torch.nn.init`` instead of ``nn.init``.\n\n### Define Model Mutations\n\nA base model is only one concrete model not a model space. We provide :doc:`API and Primitives `\nfor users to express how the base model can be mutated. That is, to build a model space which includes many models.\n\nBased on the above base model, we can define a model space as below.\n\n.. code-block:: diff\n\n @model_wrapper\n class Net(nn.Module):\n def __init__(self):\n super().__init__()\n self.conv1 = nn.Conv2d(1, 32, 3, 1)\n - self.conv2 = nn.Conv2d(32, 64, 3, 1)\n + self.conv2 = nn.LayerChoice([\n + nn.Conv2d(32, 64, 3, 1),\n + DepthwiseSeparableConv(32, 64)\n + ])\n - self.dropout1 = nn.Dropout(0.25)\n + self.dropout1 = nn.Dropout(nn.ValueChoice([0.25, 0.5, 0.75]))\n self.dropout2 = nn.Dropout(0.5)\n - self.fc1 = nn.Linear(9216, 128)\n - self.fc2 = nn.Linear(128, 10)\n + feature = nn.ValueChoice([64, 128, 256])\n + self.fc1 = nn.Linear(9216, feature)\n + self.fc2 = nn.Linear(feature, 10)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n x = F.max_pool2d(self.conv2(x), 2)\n x = torch.flatten(self.dropout1(x), 1)\n x = self.fc2(self.dropout2(F.relu(self.fc1(x))))\n output = F.log_softmax(x, dim=1)\n return output\n\nThis results in the following code:\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class DepthwiseSeparableConv(nn.Module):\n def __init__(self, in_ch, out_ch):\n super().__init__()\n self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size=3, groups=in_ch)\n self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1)\n\n def forward(self, x):\n return self.pointwise(self.depthwise(x))\n\n\n@model_wrapper\nclass ModelSpace(nn.Module):\n def __init__(self):\n super().__init__()\n self.conv1 = nn.Conv2d(1, 32, 3, 1)\n # LayerChoice is used to select a layer between Conv2d and DwConv.\n self.conv2 = nn.LayerChoice([\n nn.Conv2d(32, 64, 3, 1),\n DepthwiseSeparableConv(32, 64)\n ])\n # ValueChoice is used to select a dropout rate.\n # ValueChoice can be used as parameter of modules wrapped in `nni.retiarii.nn.pytorch`\n # or customized modules wrapped with `@basic_unit`.\n self.dropout1 = nn.Dropout(nn.ValueChoice([0.25, 0.5, 0.75])) # choose dropout rate from 0.25, 0.5 and 0.75\n self.dropout2 = nn.Dropout(0.5)\n feature = nn.ValueChoice([64, 128, 256])\n self.fc1 = nn.Linear(9216, feature)\n self.fc2 = nn.Linear(feature, 10)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n x = F.max_pool2d(self.conv2(x), 2)\n x = torch.flatten(self.dropout1(x), 1)\n x = self.fc2(self.dropout2(F.relu(self.fc1(x))))\n output = F.log_softmax(x, dim=1)\n return output\n\n\nmodel_space = ModelSpace()\nmodel_space" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This example uses two mutation APIs,\n:class:`nn.LayerChoice ` and\n:class:`nn.InputChoice `.\n:class:`nn.LayerChoice `\ntakes a list of candidate modules (two in this example), one will be chosen for each sampled model.\nIt can be used like normal PyTorch module.\n:class:`nn.InputChoice ` takes a list of candidate values,\none will be chosen to take effect for each sampled model.\n\nMore detailed API description and usage can be found :doc:`here `.\n\n

Note

We are actively enriching the mutation APIs, to facilitate easy construction of model space.\n If the currently supported mutation APIs cannot express your model space,\n please refer to :doc:`this doc ` for customizing mutators.

\n\n## Explore the Defined Model Space\n\nThere are basically two exploration approaches: (1) search by evaluating each sampled model independently,\nwhich is the search approach in `multi-trial NAS `\nand (2) one-shot weight-sharing based search, which is used in one-shot NAS.\nWe demonstrate the first approach in this tutorial. Users can refer to `here ` for the second approach.\n\nFirst, users need to pick a proper exploration strategy to explore the defined model space.\nSecond, users need to pick or customize a model evaluator to evaluate the performance of each explored model.\n\n### Pick an exploration strategy\n\nRetiarii supports many :doc:`exploration strategies `.\n\nSimply choosing (i.e., instantiate) an exploration strategy as below.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import nni.retiarii.strategy as strategy\nsearch_strategy = strategy.Random(dedup=True) # dedup=False if deduplication is not wanted" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Pick or customize a model evaluator\n\nIn the exploration process, the exploration strategy repeatedly generates new models. A model evaluator is for training\nand validating each generated model to obtain the model's performance.\nThe performance is sent to the exploration strategy for the strategy to generate better models.\n\nRetiarii has provided :doc:`built-in model evaluators `, but to start with,\nit is recommended to use :class:`FunctionalEvaluator `,\nthat is, to wrap your own training and evaluation code with one single function.\nThis function should receive one single model class and uses :func:`nni.report_final_result` to report the final score of this model.\n\nAn example here creates a simple evaluator that runs on MNIST dataset, trains for 2 epochs, and reports its validation accuracy.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import nni\n\nfrom torchvision import transforms\nfrom torchvision.datasets import MNIST\nfrom torch.utils.data import DataLoader\n\n\ndef train_epoch(model, device, train_loader, optimizer, epoch):\n loss_fn = torch.nn.CrossEntropyLoss()\n model.train()\n for batch_idx, (data, target) in enumerate(train_loader):\n data, target = data.to(device), target.to(device)\n optimizer.zero_grad()\n output = model(data)\n loss = loss_fn(output, target)\n loss.backward()\n optimizer.step()\n if batch_idx % 10 == 0:\n print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n epoch, batch_idx * len(data), len(train_loader.dataset),\n 100. * batch_idx / len(train_loader), loss.item()))\n\n\ndef test_epoch(model, device, test_loader):\n model.eval()\n test_loss = 0\n correct = 0\n with torch.no_grad():\n for data, target in test_loader:\n data, target = data.to(device), target.to(device)\n output = model(data)\n pred = output.argmax(dim=1, keepdim=True)\n correct += pred.eq(target.view_as(pred)).sum().item()\n\n test_loss /= len(test_loader.dataset)\n accuracy = 100. * correct / len(test_loader.dataset)\n\n print('\\nTest set: Accuracy: {}/{} ({:.0f}%)\\n'.format(\n correct, len(test_loader.dataset), accuracy))\n\n return accuracy\n\n\ndef evaluate_model(model_cls):\n # \"model_cls\" is a class, need to instantiate\n model = model_cls()\n\n device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n model.to(device)\n\n optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n transf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n train_loader = DataLoader(MNIST('data/mnist', download=True, transform=transf), batch_size=64, shuffle=True)\n test_loader = DataLoader(MNIST('data/mnist', download=True, train=False, transform=transf), batch_size=64)\n\n for epoch in range(3):\n # train the model for one epoch\n train_epoch(model, device, train_loader, optimizer, epoch)\n # test the model for one epoch\n accuracy = test_epoch(model, device, test_loader)\n # call report intermediate result. Result can be float or dict\n nni.report_intermediate_result(accuracy)\n\n # report final test result\n nni.report_final_result(accuracy)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Create the evaluator\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from nni.retiarii.evaluator import FunctionalEvaluator\nevaluator = FunctionalEvaluator(evaluate_model)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The ``train_epoch`` and ``test_epoch`` here can be any customized function,\nwhere users can write their own training recipe.\n\nIt is recommended that the ``evaluate_model`` here accepts no additional arguments other than ``model_cls``.\nHowever, in the :doc:`advanced tutorial `, we will show how to use additional arguments in case you actually need those.\nIn future, we will support mutation on the arguments of evaluators, which is commonly called \"Hyper-parmeter tuning\".\n\n## Launch an Experiment\n\nAfter all the above are prepared, it is time to start an experiment to do the model search. An example is shown below.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig\nexp = RetiariiExperiment(model_space, evaluator, [], search_strategy)\nexp_config = RetiariiExeConfig('local')\nexp_config.experiment_name = 'mnist_search'" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The following configurations are useful to control how many trials to run at most / at the same time.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "exp_config.max_trial_number = 4 # spawn 4 trials at most\nexp_config.trial_concurrency = 2 # will run two trials concurrently" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Remember to set the following config if you want to GPU.\n``use_active_gpu`` should be set true if you wish to use an occupied GPU (possibly running a GUI).\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "exp_config.trial_gpu_number = 1\nexp_config.training_service.use_active_gpu = True" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Launch the experiment. The experiment should take several minutes to finish on a workstation with 2 GPUs.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "exp.run(exp_config, 8081)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Users can also run Retiarii Experiment with :doc:`different training services `\nbesides ``local`` training service.\n\n## Visualize the Experiment\n\nUsers can visualize their experiment in the same way as visualizing a normal hyper-parameter tuning experiment.\nFor example, open ``localhost:8081`` in your browser, 8081 is the port that you set in ``exp.run``.\nPlease refer to :doc:`here ` for details.\n\nWe support visualizing models with 3rd-party visualization engines (like `Netron `__).\nThis can be used by clicking ``Visualization`` in detail panel for each trial.\nNote that current visualization is based on `onnx `__ ,\nthus visualization is not feasible if the model cannot be exported into onnx.\n\nBuilt-in evaluators (e.g., Classification) will automatically export the model into a file.\nFor your own evaluator, you need to save your file into ``$NNI_OUTPUT_DIR/model.onnx`` to make this work.\nFor instance,\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import os\nfrom pathlib import Path\n\n\ndef evaluate_model_with_visualization(model_cls):\n model = model_cls()\n # dump the model into an onnx\n if 'NNI_OUTPUT_DIR' in os.environ:\n dummy_input = torch.zeros(1, 3, 32, 32)\n torch.onnx.export(model, (dummy_input, ),\n Path(os.environ['NNI_OUTPUT_DIR']) / 'model.onnx')\n evaluate_model(model_cls)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Relaunch the experiment, and a button is shown on Web portal.\n\n\n\n## Export Top Models\n\nUsers can export top models after the exploration is done using ``export_top_models``.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "for model_dict in exp.export_top_models(formatter='dict'):\n print(model_dict)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The output is ``json`` object which records the mutation actions of the top model.\nIf users want to output source code of the top model,\nthey can use `graph-based execution engine ` for the experiment,\nby simply adding the following two lines.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "exp_config.execution_engine = 'base'\nexport_formatter = 'code'" ] } ], "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": 0 }