Skip to content

GitLab

Unverified Commit bddaf36d authored by Yuge Zhang's avatar Yuge Zhang Committed by GitHub
Browse files

Tutorial of searching in DARTS space (#5053)

parent cbc6273a
Show whitespace changes
Inline Side-by-side
Showing with 3175 additions and 5 deletions
docs/source/nas/toctree.rst
View file @ bddaf36d
......@@ -5,7 +5,7 @@ Neural Architecture Search
:hidden:
overview
Quickstart </tutorials/hello_nas>
Tutorials <tutorials>
construct_space
exploration_strategy
evaluator
......
docs/source/nas/tutorials.rst 0 → 100644
View file @ bddaf36d
NAS Tutorials
=============
.. toctree::
:hidden:
Hello NAS! </tutorials/hello_nas>
Search in DARTS </tutorials/darts>
docs/source/tutorials/darts.ipynb 0 → 100644
View file @ bddaf36d
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"%matplotlib inline"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n# Searching in DARTS search space\n\nIn this tutorial, we demonstrate how to search in the famous model space proposed in `DARTS`_.\n\nThrough this process, you will learn:\n\n* How to use the built-in model spaces from NNI's model space hub.\n* How to use one-shot exploration strategies to explore a model space.\n* How to customize evaluators to achieve the best performance.\n\nIn the end, we get a strong-performing model on CIFAR-10 dataset, which achieves up to 97.28% accuracy.\n\n.. attention::\n\n Running this tutorial requires a GPU.\n If you don't have one, you can set ``gpus`` in :class:`~nni.retiarii.evaluator.pytorch.Classification` to be 0,\n but do note that it will be much slower.\n\n\n## Use a pre-searched model\n\nSimilar to [the beginner tutorial of PyTorch](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html)_,\nwe begin with CIFAR-10 dataset, which is a image classification dataset of 10 categories.\nThe images in CIFAR-10 are of size 3x32x32, i.e., RGB-colored images of 32x32 pixels in size.\n\nWe first load the CIFAR-10 dataset with torchvision.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import nni\nimport torch\nfrom torchvision import transforms\nfrom torchvision.datasets import CIFAR10\nfrom nni.retiarii.evaluator.pytorch import DataLoader\n\nCIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]\nCIFAR_STD = [0.24703233, 0.24348505, 0.26158768]\n\ntransform_valid = transforms.Compose([\n transforms.ToTensor(),\n transforms.Normalize(CIFAR_MEAN, CIFAR_STD),\n])\nvalid_data = nni.trace(CIFAR10)(root='./data', train=False, download=True, transform=transform_valid)\nvalid_loader = DataLoader(valid_data, batch_size=256, num_workers=6)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<div class=\"alert alert-info\"><h4>Note</h4><p>If you are to use multi-trial strategies, wrapping CIFAR10 with :func:`nni.trace` and\n use DataLoader from ``nni.retiarii.evaluator.pytorch`` (instead of ``torch.utils.data``) are mandatory.\n Otherwise, it's optional.</p></div>\n\nWhen working with famous datasets like CIFAR-10 or ImageNet,\nit's tempting to use or finetune from a pretrained model, like ResNet.\nThere's nothing wrong with doing so, and sometimes it might be beneficial.\nThanks to the development of NAS, we now have quite a large number of *pre-searched models*,\nwhich are produced by most popular NAS literatures.\nYou can easily load these models, validate their performances, and finetune them if you need.\n\nWe present :doc:`model space hub </nas/space_hub>`, where you can find many built-in model spaces,\nalong with many pre-searched models.\nWe choose one from `DARTS`_ search space, which is natively trained on our target dataset, CIFAR-10,\nso as to save the tedious steps of finetuning.\n\n.. tip::\n\n Finetuning a pre-searched model on other datasets is no different from finetuning *any model*.\n We recommend reading\n [this tutorial of object detection finetuning](https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html)_\n if you want to know how finetuning is generally done in PyTorch.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"from nni.retiarii.hub.pytorch import DARTS as DartsSpace\n\ndarts_v2_model = DartsSpace.load_searched_model('darts-v2', pretrained=True, download=True)\n\ndef evaluate_model(model, cuda=False):\n device = torch.device('cuda' if cuda else 'cpu')\n model.to(device)\n model.eval()\n with torch.no_grad():\n correct = total = 0\n for inputs, targets in valid_loader:\n inputs, targets = inputs.to(device), targets.to(device)\n logits = model(inputs)\n _, predict = torch.max(logits, 1)\n correct += (predict == targets).sum().cpu().item()\n total += targets.size(0)\n print('Accuracy:', correct / total)\n return correct / total\n\nevaluate_model(darts_v2_model, True) # Set this to false if there's no GPU."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The journey could end here. Or you are interested,\nwe can go a step further to search a model within :class:`~nni.retiarii.hub.pytorch.DARTS` space on our own.\n\n## Use the model space\n\nThe model space provided in `DARTS`_ originated from [NASNet](https://arxiv.org/abs/1707.07012)_,\nwhere the full model is constructed by repeatedly stacking a single computational unit (called a **cell**).\nThere are two types of cells within a network. The first type is called *normal cell*, and the second type is called *reduction cell*.\nThe key difference between normal and reduction cell is that the reduction cell will downsample the input feature map,\nand decrease its resolution. Normal and reduction cells are stacked alternately, as shown in the following figure.\n\n<img src=\"file://../../img/nasnet_cell_stack.png\">\n\nA cell takes outputs from two previous cells as inputs and contains a collection of *nodes*.\nEach node takes two previous nodes within the same cell (or the two cell inputs),\nand applies an *operator* (e.g., convolution, or max-pooling) to each input,\nand sums the outputs of operators as the output of the node.\nThe output of cell is the concatenation of all the nodes that are never used as inputs of another node.\nWe recommend reading [NDS](https://arxiv.org/pdf/1905.13214.pdf)_ or [ENAS](https://arxiv.org/abs/1802.03268)_ for details.\n\nWe illustrate an example of cells in the following figure.\n\n<img src=\"file://../../img/nasnet_cell.png\">\n\nThe search space proposed in `DARTS`_ paper introduced two modifications to the original space\nin [NASNet](https://arxiv.org/abs/1707.07012)_.\n\nFirstly, the operator candidates have been narrowed down to seven:\n\n- Max pooling 3x3\n- Average pooling 3x3\n- Skip connect (Identity)\n- Separable convolution 3x3\n- Separable convolution 5x5\n- Dilated convolution 3x3\n- Dilated convolution 5x5\n\nSecondly, the output of cell is the concatenate of **all the nodes within the cell**.\n\nAs the search space is based on cell, once the normal and reduction cell has been fixed, we can stack them for indefinite times.\nTo save the search cost, the common practice is to reduce the number of filters (i.e., channels) and number of stacked cells\nduring the search phase, and increase them back when training the final searched architecture.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>`DARTS`_ is one of those papers that innovate both in search space and search strategy.\n In this tutorial, we will search on **model space** provided by DARTS with **search strategy** proposed by DARTS.\n We refer to them as *DARTS model space* (``DartsSpace``) and *DARTS strategy* (``DartsStrategy``), respectively.\n We did NOT imply that the :class:`~nni.retiarii.hub.pytorch.DARTS` space and\n :class:`~nni.retiarii.strategy.DARTS` strategy has to used together.\n You can always explore the DARTS space with another search strategy, or use your own strategy to search a different model space.</p></div>\n\nIn the following example, we initialize a :class:`~nni.retiarii.hub.pytorch.DARTS`\nmodel space, with 16 initial filters and 8 stacked cells.\nThe network is specialized for CIFAR-10 dataset with 32x32 input resolution.\n\nThe :class:`~nni.retiarii.hub.pytorch.DARTS` model space here is provided by :doc:`model space hub </nas/space_hub>`,\nwhere we have supported multiple popular model spaces for plug-and-play.\n\n.. tip::\n\n The model space here can be replaced with any space provided in the hub,\n or even customized spaces built from scratch.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"model_space = DartsSpace(16, 8, 'cifar')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Search on the model space\n\n<div class=\"alert alert-danger\"><h4>Warning</h4><p>Please set ``fast_dev_run`` to False to reproduce the our claimed results.\n Otherwise, only a few mini-batches will be run.</p></div>\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"fast_dev_run = True"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Evaluator\n\nTo begin exploring the model space, one firstly need to have an evaluator to provide the criterion of a \"good model\".\nAs we are searching on CIFAR-10 dataset, one can easily use the :class:`~nni.retiarii.evaluator.pytorch.Classification`\nas a starting point.\n\nNote that for a typical setup of NAS, the model search should be on validation set, and the evaluation of the final searched model\nshould be on test set. However, as CIFAR-10 dataset doesn't have a test dataset (only 50k train + 10k valid),\nwe have to split the original training set into a training set and a validation set.\nThe recommended train/val split by `DARTS`_ strategy is 1:1.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import numpy as np\nfrom nni.retiarii.evaluator.pytorch import Classification\nfrom torch.utils.data import SubsetRandomSampler\n\ntransform = transforms.Compose([\n transforms.RandomCrop(32, padding=4),\n transforms.RandomHorizontalFlip(),\n transforms.ToTensor(),\n transforms.Normalize(CIFAR_MEAN, CIFAR_STD),\n])\n\ntrain_data = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform)\n\nnum_samples = len(train_data)\nindices = np.random.permutation(num_samples)\nsplit = num_samples // 2\n\nsearch_train_loader = DataLoader(\n train_data, batch_size=64, num_workers=6,\n sampler=SubsetRandomSampler(indices[:split]),\n)\n\nsearch_valid_loader = DataLoader(\n train_data, batch_size=64, num_workers=6,\n sampler=SubsetRandomSampler(indices[split:]),\n)\n\nevaluator = Classification(\n learning_rate=1e-3,\n weight_decay=1e-4,\n train_dataloaders=search_train_loader,\n val_dataloaders=search_valid_loader,\n max_epochs=10,\n gpus=1,\n fast_dev_run=fast_dev_run,\n)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Strategy\n\nWe will use `DARTS`_ (Differentiable ARchiTecture Search) as the search strategy to explore the model space.\n:class:`~nni.retiarii.strategy.DARTS` strategy belongs to the category of `one-shot strategy <one-shot-nas>`.\nThe fundamental differences between One-shot strategies and `multi-trial strategies <multi-trial-nas>` is that,\none-shot strategy combines search with model training into a single run.\nCompared to multi-trial strategies, one-shot NAS doesn't need to iteratively spawn new trials (i.e., models),\nand thus saves the excessive cost of model training.\nIt's worth mentioning that one-shot NAS also suffers from multiple drawbacks despite its computational efficiency.\nWe recommend\n[Weight-Sharing Neural Architecture Search: A Battle to Shrink the Optimization Gap](https://arxiv.org/abs/2008.01475)_\nand\n[How Does Supernet Help in Neural Architecture Search?](https://arxiv.org/abs/2010.08219)_ for interested readers.\n\n:class:`~nni.retiarii.strategy.DARTS` strategy is provided as one of NNI's :doc:`built-in search strategies </nas/exploration_strategy>`.\nUsing it can be as simple as one line of code.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"from nni.retiarii.strategy import DARTS as DartsStrategy\n\nstrategy = DartsStrategy()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
".. tip:: The ``DartsStrategy`` here can be replaced by any search strategies, even multi-trial strategies.\n\nIf you want to know how DARTS strategy works, here is a brief version.\nUnder the hood, DARTS converts the cell into a densely connected graph, and put operators on edges (see the following figure).\nSince the operators are not decided yet, every edge is a weighted mixture of multiple operators (multiple color in the figure).\nDARTS then learns to assign the optimal \"color\" for each edge during the network training.\nIt finally selects one \"color\" for each edge, and drops redundant edges.\nThe weights on the edges are called *architecture weights*.\n\n<img src=\"file://../../img/darts_illustration.png\">\n\nIt's NOT reflected in the figure that, for DARTS model space, exactly two inputs are kept for every node.\n\n### Launch experiment\n\nWe then come to the step of launching the experiment.\nThis step is similar to what we have done in the :doc:`beginner tutorial <hello_nas>`,\nexcept that the ``execution_engine`` argument should be set to ``oneshot``.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig\n\nconfig = RetiariiExeConfig(execution_engine='oneshot')\nexperiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)\nexperiment.run(config)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
".. tip::\n\n The search process can be visualized with tensorboard. For example::\n\n tensorboard --logdir=./lightning_logs\n\n Then, open the browser and go to http://localhost:6006/ to monitor the search process.\n\n .. image:: ../../img/darts_search_process.png\n\nWe can then retrieve the best model found by the strategy with ``export_top_models``.\nHere, the retrieved model is a dict (called *architecture dict*) describing the selected normal cell and reduction cell.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"exported_arch = experiment.export_top_models()[0]\n\nexported_arch"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The cell can be visualized with the following code snippet\n(copied and modified from [DARTS visualization](https://github.com/quark0/darts/blob/master/cnn/visualize.py)_).\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import io\nimport graphviz\nimport matplotlib.pyplot as plt\nfrom PIL import Image\n\ndef plot_single_cell(arch_dict, cell_name):\n g = graphviz.Digraph(\n node_attr=dict(style='filled', shape='rect', align='center'),\n format='png'\n )\n g.body.extend(['rankdir=LR'])\n\n g.node('c_{k-2}', fillcolor='darkseagreen2')\n g.node('c_{k-1}', fillcolor='darkseagreen2')\n assert len(arch_dict) % 2 == 0\n\n for i in range(2, 6):\n g.node(str(i), fillcolor='lightblue')\n\n for i in range(2, 6):\n for j in range(2):\n op = arch_dict[f'{cell_name}/op_{i}_{j}']\n from_ = arch_dict[f'{cell_name}/input_{i}_{j}']\n if from_ == 0:\n u = 'c_{k-2}'\n elif from_ == 1:\n u = 'c_{k-1}'\n else:\n u = str(from_)\n v = str(i)\n g.edge(u, v, label=op, fillcolor='gray')\n\n g.node('c_{k}', fillcolor='palegoldenrod')\n for i in range(2, 6):\n g.edge(str(i), 'c_{k}', fillcolor='gray')\n\n g.attr(label=f'{cell_name.capitalize()} cell')\n\n image = Image.open(io.BytesIO(g.pipe()))\n return image\n\ndef plot_double_cells(arch_dict):\n image1 = plot_single_cell(arch_dict, 'normal')\n image2 = plot_single_cell(arch_dict, 'reduce')\n height_ratio = max(image1.size[1] / image1.size[0], image2.size[1] / image2.size[0]) \n _, axs = plt.subplots(1, 2, figsize=(20, 10 * height_ratio))\n axs[0].imshow(image1)\n axs[1].imshow(image2)\n axs[0].axis('off')\n axs[1].axis('off')\n plt.show()\n\nplot_double_cells(exported_arch)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<div class=\"alert alert-danger\"><h4>Warning</h4><p>The cell above is obtained via ``fast_dev_run`` (i.e., running only 1 mini-batch).</p></div>\n\nWhen ``fast_dev_run`` is turned off, we get a model with the following architecture,\nwhere you might notice an interesting fact that around half the operations have selected ``sep_conv_3x3``.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"plot_double_cells({\n 'normal/op_2_0': 'sep_conv_3x3',\n 'normal/input_2_0': 1,\n 'normal/op_2_1': 'sep_conv_3x3',\n 'normal/input_2_1': 0,\n 'normal/op_3_0': 'sep_conv_3x3',\n 'normal/input_3_0': 1,\n 'normal/op_3_1': 'sep_conv_3x3',\n 'normal/input_3_1': 2,\n 'normal/op_4_0': 'sep_conv_3x3',\n 'normal/input_4_0': 1,\n 'normal/op_4_1': 'sep_conv_3x3',\n 'normal/input_4_1': 0,\n 'normal/op_5_0': 'sep_conv_3x3',\n 'normal/input_5_0': 1,\n 'normal/op_5_1': 'max_pool_3x3',\n 'normal/input_5_1': 0,\n 'reduce/op_2_0': 'sep_conv_3x3',\n 'reduce/input_2_0': 0,\n 'reduce/op_2_1': 'sep_conv_3x3',\n 'reduce/input_2_1': 1,\n 'reduce/op_3_0': 'dil_conv_5x5',\n 'reduce/input_3_0': 2,\n 'reduce/op_3_1': 'sep_conv_3x3',\n 'reduce/input_3_1': 0,\n 'reduce/op_4_0': 'dil_conv_5x5',\n 'reduce/input_4_0': 2,\n 'reduce/op_4_1': 'sep_conv_5x5',\n 'reduce/input_4_1': 1,\n 'reduce/op_5_0': 'sep_conv_5x5',\n 'reduce/input_5_0': 4,\n 'reduce/op_5_1': 'dil_conv_5x5',\n 'reduce/input_5_1': 2\n})"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Retrain the searched model\n\nWhat we have got in the last step, is only a cell structure.\nTo get a final usable model with trained weights, we need to construct a real model based on this structure,\nand then fully train it.\n\nTo construct a fixed model based on the architecture dict exported from the experiment,\nwe can use :func:`nni.retiarii.fixed_arch`. Seemingly, we are still creating a space.\nBut under the with-context, we are actually creating a fixed model.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"from nni.retiarii import fixed_arch\n\nwith fixed_arch(exported_arch):\n final_model = DartsSpace(16, 8, 'cifar')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We then train the model on full CIFAR-10 training dataset, and evaluate it on the original CIFAR-10 validation dataset.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"train_loader = DataLoader(train_data, batch_size=96, num_workers=6) # Use the original training data"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The validation data loader can be reused.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"valid_loader"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We must create a new evaluator here because a different data split is used.\nAlso, we should avoid the underlying pytorch-lightning implementation of :class:`~nni.retiarii.evaluator.pytorch.Classification`\nevaluator from loading the wrong checkpoint.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"max_epochs = 100\n\nevaluator = Classification(\n learning_rate=1e-3,\n weight_decay=1e-4,\n train_dataloaders=train_loader,\n val_dataloaders=valid_loader,\n max_epochs=max_epochs,\n gpus=1,\n export_onnx=False, # Disable ONNX export for this experiment\n fast_dev_run=fast_dev_run, # Should be false for fully training\n)\n\nevaluator.fit(final_model)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<div class=\"alert alert-info\"><h4>Note</h4><p>When ``fast_dev_run`` is turned off, we achieve a validation accuracy of 89.69% after training for 100 epochs.</p></div>\n\n## Reproduce results in DARTS paper\n\nAfter a brief walkthrough of search + retrain process with one-shot strategy,\nwe then fill the gap between our results (89.69%) and the results in the `DARTS` paper.\nThis is because we didn't introduce some extra training tricks, including [DropPath](https://arxiv.org/pdf/1605.07648v4.pdf)_,\nAuxiliary loss, gradient clipping and augmentations like [Cutout](https://arxiv.org/pdf/1708.04552v2.pdf)_.\nThey also train the deeper (20 cells) and wider (36 channels) networks for longer time (600 epochs).\n\n\n### Evaluator\n\nTo implement these tricks, we first need to rewrite a few parts of evaluator.\n\nWorking with one-shot strategies, evaluators need to be implemented in the style of `PyTorch-Lightning <lightning-evaluator>`,\nThe full tutorial can be found in :doc:`/nas/evaluator`.\nPutting it briefly, the core part of writing a new evaluator is to write a new LightningModule.\n[LightingModule](https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html)_ is a concept in\nPyTorch-Lightning, which organizes the model training process into a list of functions, such as,\n``training_step``, ``validation_step``, ``configure_optimizers``, etc.\nSince we are merely adding a few ingredients to :class:`~nni.retiarii.evaluator.pytorch.Classification`,\nwe can simply inherit :class:`~nni.retiarii.evaluator.pytorch.ClassificationModule`, which is the underlying LightningModule\nbehind :class:`~nni.retiarii.evaluator.pytorch.Classification`.\nThis could look intimidating at first, but most of them are just plug-and-play tricks which you don't need to know details about.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import torch\nfrom nni.retiarii.evaluator.pytorch import ClassificationModule\n\nclass DartsClassificationModule(ClassificationModule):\n def __init__(\n self,\n learning_rate: float = 0.001,\n weight_decay: float = 0.,\n auxiliary_loss_weight: float = 0.4,\n max_epochs: int = 600\n ):\n self.auxiliary_loss_weight = auxiliary_loss_weight\n # Training length will be used in LR scheduler\n self.max_epochs = max_epochs\n super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False)\n\n def configure_optimizers(self):\n \"\"\"Customized optimizer with momentum, as well as a scheduler.\"\"\"\n optimizer = torch.optim.SGD(\n self.parameters(),\n momentum=0.9,\n lr=self.hparams.learning_rate,\n weight_decay=self.hparams.weight_decay\n )\n return {\n 'optimizer': optimizer,\n 'lr_scheduler': torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, self.max_epochs, eta_min=1e-3)\n }\n\n def training_step(self, batch, batch_idx):\n \"\"\"Training step, customized with auxiliary loss.\"\"\"\n x, y = batch\n if self.auxiliary_loss_weight:\n y_hat, y_aux = self(x)\n loss_main = self.criterion(y_hat, y)\n loss_aux = self.criterion(y_aux, y)\n self.log('train_loss_main', loss_main)\n self.log('train_loss_aux', loss_aux)\n loss = loss_main + self.auxiliary_loss_weight * loss_aux\n else:\n y_hat = self(x)\n loss = self.criterion(y_hat, y)\n self.log('train_loss', loss, prog_bar=True)\n for name, metric in self.metrics.items():\n self.log('train_' + name, metric(y_hat, y), prog_bar=True)\n return loss\n\n def on_train_epoch_start(self):\n # Set drop path probability before every epoch. This has no effect if drop path is not enabled in model.\n self.model.set_drop_path_prob(self.model.drop_path_prob * self.current_epoch / self.max_epochs)\n\n # Logging learning rate at the beginning of every epoch\n self.log('lr', self.trainer.optimizers[0].param_groups[0]['lr'])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The full evaluator is written as follows,\nwhich simply wraps everything (except model space and search strategy of course), in a single object.\n:class:`~nni.retiarii.evaluator.pytorch.Lightning` here is a special type of evaluator.\nDon't forget to use the train/val data split specialized for search (1:1) here.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"from nni.retiarii.evaluator.pytorch import Lightning, Trainer\n\nmax_epochs = 50\n\nevaluator = Lightning(\n DartsClassificationModule(0.025, 3e-4, 0., max_epochs),\n Trainer(\n gpus=1,\n max_epochs=max_epochs,\n fast_dev_run=fast_dev_run,\n ),\n train_dataloaders=search_train_loader,\n val_dataloaders=search_valid_loader\n)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Strategy\n\n:class:`~nni.retiarii.strategy.DARTS` strategy is created with gradient clip turned on.\nIf you are familiar with PyTorch-Lightning, you might aware that gradient clipping can be enabled in Lightning trainer.\nHowever, enabling gradient cip in the trainer above won't work, because the underlying\nimplementation of :class:`~nni.retiarii.strategy.DARTS` strategy is based on\n[manual optimization](https://pytorch-lightning.readthedocs.io/en/stable/common/optimization.html)_.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"strategy = DartsStrategy(gradient_clip_val=5.)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Launch experiment\n\nThen we use the newly created evaluator and strategy to launch the experiment again.\n\n<div class=\"alert alert-danger\"><h4>Warning</h4><p>``model_space`` has to be re-instantiated because a known limitation,\n i.e., one model space can't be reused across multiple experiments.</p></div>\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"model_space = DartsSpace(16, 8, 'cifar')\n\nconfig = RetiariiExeConfig(execution_engine='oneshot')\nexperiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)\nexperiment.run(config)\n\nexported_arch = experiment.export_top_models()[0]\n\nexported_arch"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We get the following architecture when ``fast_dev_run`` is set to False. It takes around 8 hours on a P100 GPU.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"plot_double_cells({\n 'normal/op_2_0': 'sep_conv_3x3',\n 'normal/input_2_0': 0,\n 'normal/op_2_1': 'sep_conv_3x3',\n 'normal/input_2_1': 1,\n 'normal/op_3_0': 'sep_conv_3x3',\n 'normal/input_3_0': 1,\n 'normal/op_3_1': 'skip_connect',\n 'normal/input_3_1': 0,\n 'normal/op_4_0': 'sep_conv_3x3',\n 'normal/input_4_0': 0,\n 'normal/op_4_1': 'max_pool_3x3',\n 'normal/input_4_1': 1,\n 'normal/op_5_0': 'sep_conv_3x3',\n 'normal/input_5_0': 0,\n 'normal/op_5_1': 'sep_conv_3x3',\n 'normal/input_5_1': 1,\n 'reduce/op_2_0': 'max_pool_3x3',\n 'reduce/input_2_0': 0,\n 'reduce/op_2_1': 'sep_conv_5x5',\n 'reduce/input_2_1': 1,\n 'reduce/op_3_0': 'dil_conv_5x5',\n 'reduce/input_3_0': 2,\n 'reduce/op_3_1': 'max_pool_3x3',\n 'reduce/input_3_1': 0,\n 'reduce/op_4_0': 'max_pool_3x3',\n 'reduce/input_4_0': 0,\n 'reduce/op_4_1': 'sep_conv_3x3',\n 'reduce/input_4_1': 2,\n 'reduce/op_5_0': 'max_pool_3x3',\n 'reduce/input_5_0': 0,\n 'reduce/op_5_1': 'skip_connect',\n 'reduce/input_5_1': 2\n})"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Retrain\n\nWhen retraining,\nwe extend the original dataloader to introduce another trick called [Cutout](https://arxiv.org/pdf/1708.04552v2.pdf)_.\nCutout is a data augmentation technique that randomly masks out rectangular regions in images.\nIn CIFAR-10, the typical masked size is 16x16 (the image sizes are 32x32 in the dataset).\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"def cutout_transform(img, length: int = 16):\n h, w = img.size(1), img.size(2)\n mask = np.ones((h, w), np.float32)\n y = np.random.randint(h)\n x = np.random.randint(w)\n\n y1 = np.clip(y - length // 2, 0, h)\n y2 = np.clip(y + length // 2, 0, h)\n x1 = np.clip(x - length // 2, 0, w)\n x2 = np.clip(x + length // 2, 0, w)\n\n mask[y1: y2, x1: x2] = 0.\n mask = torch.from_numpy(mask)\n mask = mask.expand_as(img)\n img *= mask\n return img\n\ntransform_with_cutout = transforms.Compose([\n transforms.RandomCrop(32, padding=4),\n transforms.RandomHorizontalFlip(),\n transforms.ToTensor(),\n transforms.Normalize(CIFAR_MEAN, CIFAR_STD),\n cutout_transform,\n])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The train dataloader needs to be reinstantiated with the new transform.\nThe validation dataloader is not affected, and thus can be reused.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"train_data_cutout = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform_with_cutout)\ntrain_loader_cutout = DataLoader(train_data_cutout, batch_size=96)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We then create the final model based on the new exported architecture.\nThis time, auxiliary loss and drop path probability is enabled.\n\nFollowing the same procedure as paper, we also increase the number of filters to 36, and number of cells to 20,\nso as to reasonably increase the model size and boost the performance.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"with fixed_arch(exported_arch):\n final_model = DartsSpace(36, 20, 'cifar', auxiliary_loss=True, drop_path_prob=0.2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We create a new evaluator for the retraining process, where the gradient clipping is put into the keyword arguments of trainer.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"max_epochs = 600\n\nevaluator = Lightning(\n DartsClassificationModule(0.025, 3e-4, 0.4, max_epochs),\n Trainer(\n gpus=1,\n gradient_clip_val=5.,\n max_epochs=max_epochs,\n fast_dev_run=fast_dev_run\n ),\n train_dataloaders=train_loader_cutout,\n val_dataloaders=valid_loader,\n)\n\nevaluator.fit(final_model)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When ``fast_dev_run`` is turned off, after retraining, the architecture yields a top-1 accuracy of 97.12%.\nIf we take the best snapshot throughout the retrain process,\nthere is a chance that the top-1 accuracy will be 97.28%.\n\n<img src=\"file://../../img/darts_val_acc.png\">\n\nIn the figure, the orange line is the validation accuracy curve after training for 600 epochs,\nwhile the red line corresponding the previous version in this tutorial before adding all the training tricks and\nonly trains for 100 epochs.\n\nThe results outperforms \"DARTS (first order) + cutout\" in `DARTS`_ paper, which is only 97.00\u00b10.14%.\nIt's even comparable with \"DARTS (second order) + cutout\" in the paper (97.24\u00b10.09%),\nthough we didn't implement the second order version.\nThe implementation of second order DARTS is in our future plan, and we also welcome your contribution.\n\n"
]
}
],
"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.13"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
\ No newline at end of file
docs/source/tutorials/darts.py 0 → 100644
View file @ bddaf36d
"""
Searching in DARTS search space
===============================
In this tutorial, we demonstrate how to search in the famous model space proposed in `DARTS`_.
Through this process, you will learn:
* How to use the built-in model spaces from NNI's model space hub.
* How to use one-shot exploration strategies to explore a model space.
* How to customize evaluators to achieve the best performance.
In the end, we get a strong-performing model on CIFAR-10 dataset, which achieves up to 97.28% accuracy.
.. attention::
Running this tutorial requires a GPU.
If you don't have one, you can set ``gpus`` in :class:`~nni.retiarii.evaluator.pytorch.Classification` to be 0,
but do note that it will be much slower.
.. _DARTS: https://arxiv.org/abs/1806.09055
Use a pre-searched model
------------------------
Similar to `the beginner tutorial of PyTorch <https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html>`__,
we begin with CIFAR-10 dataset, which is a image classification dataset of 10 categories.
The images in CIFAR-10 are of size 3x32x32, i.e., RGB-colored images of 32x32 pixels in size.
We first load the CIFAR-10 dataset with torchvision.
"""
import nni
import torch
from torchvision import transforms
from torchvision.datasets import CIFAR10
from nni.retiarii.evaluator.pytorch import DataLoader
CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]
transform_valid = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
valid_data = nni.trace(CIFAR10)(root='./data', train=False, download=True, transform=transform_valid)
valid_loader = DataLoader(valid_data, batch_size=256, num_workers=6)
# %%
#
# .. note::
#
# If you are to use multi-trial strategies, wrapping CIFAR10 with :func:`nni.trace` and
# use DataLoader from ``nni.retiarii.evaluator.pytorch`` (instead of ``torch.utils.data``) are mandatory.
# Otherwise, it's optional.
#
# When working with famous datasets like CIFAR-10 or ImageNet,
# it's tempting to use or finetune from a pretrained model, like ResNet.
# There's nothing wrong with doing so, and sometimes it might be beneficial.
# Thanks to the development of NAS, we now have quite a large number of *pre-searched models*,
# which are produced by most popular NAS literatures.
# You can easily load these models, validate their performances, and finetune them if you need.
#
# We present :doc:`model space hub </nas/space_hub>`, where you can find many built-in model spaces,
# along with many pre-searched models.
# We choose one from `DARTS`_ search space, which is natively trained on our target dataset, CIFAR-10,
# so as to save the tedious steps of finetuning.
#
# .. tip::
#
# Finetuning a pre-searched model on other datasets is no different from finetuning *any model*.
# We recommend reading
# `this tutorial of object detection finetuning <https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html>`__
# if you want to know how finetuning is generally done in PyTorch.
from nni.retiarii.hub.pytorch import DARTS as DartsSpace
darts_v2_model = DartsSpace.load_searched_model('darts-v2', pretrained=True, download=True)
def evaluate_model(model, cuda=False):
device = torch.device('cuda' if cuda else 'cpu')
model.to(device)
model.eval()
with torch.no_grad():
correct = total = 0
for inputs, targets in valid_loader:
inputs, targets = inputs.to(device), targets.to(device)
logits = model(inputs)
_, predict = torch.max(logits, 1)
correct += (predict == targets).sum().cpu().item()
total += targets.size(0)
print('Accuracy:', correct / total)
return correct / total
evaluate_model(darts_v2_model, True) # Set this to false if there's no GPU.
# %%
#
# The journey could end here. Or you are interested,
# we can go a step further to search a model within :class:`~nni.retiarii.hub.pytorch.DARTS` space on our own.
#
# Use the model space
# -------------------
#
# The model space provided in `DARTS`_ originated from `NASNet <https://arxiv.org/abs/1707.07012>`__,
# where the full model is constructed by repeatedly stacking a single computational unit (called a **cell**).
# There are two types of cells within a network. The first type is called *normal cell*, and the second type is called *reduction cell*.
# The key difference between normal and reduction cell is that the reduction cell will downsample the input feature map,
# and decrease its resolution. Normal and reduction cells are stacked alternately, as shown in the following figure.
#
# .. image:: ../../img/nasnet_cell_stack.png
#
# A cell takes outputs from two previous cells as inputs and contains a collection of *nodes*.
# Each node takes two previous nodes within the same cell (or the two cell inputs),
# and applies an *operator* (e.g., convolution, or max-pooling) to each input,
# and sums the outputs of operators as the output of the node.
# The output of cell is the concatenation of all the nodes that are never used as inputs of another node.
# We recommend reading `NDS <https://arxiv.org/pdf/1905.13214.pdf>`__ or `ENAS <https://arxiv.org/abs/1802.03268>`__ for details.
#
# We illustrate an example of cells in the following figure.
#
# .. image:: ../../img/nasnet_cell.png
#
# The search space proposed in `DARTS`_ paper introduced two modifications to the original space
# in `NASNet <https://arxiv.org/abs/1707.07012>`__.
#
# Firstly, the operator candidates have been narrowed down to seven:
#
# - Max pooling 3x3
# - Average pooling 3x3
# - Skip connect (Identity)
# - Separable convolution 3x3
# - Separable convolution 5x5
# - Dilated convolution 3x3
# - Dilated convolution 5x5
#
# Secondly, the output of cell is the concatenate of **all the nodes within the cell**.
#
# As the search space is based on cell, once the normal and reduction cell has been fixed, we can stack them for indefinite times.
# To save the search cost, the common practice is to reduce the number of filters (i.e., channels) and number of stacked cells
# during the search phase, and increase them back when training the final searched architecture.
#
# .. note::
#
# `DARTS`_ is one of those papers that innovate both in search space and search strategy.
# In this tutorial, we will search on **model space** provided by DARTS with **search strategy** proposed by DARTS.
# We refer to them as *DARTS model space* (``DartsSpace``) and *DARTS strategy* (``DartsStrategy``), respectively.
# We did NOT imply that the :class:`~nni.retiarii.hub.pytorch.DARTS` space and
# :class:`~nni.retiarii.strategy.DARTS` strategy has to used together.
# You can always explore the DARTS space with another search strategy, or use your own strategy to search a different model space.
#
# In the following example, we initialize a :class:`~nni.retiarii.hub.pytorch.DARTS`
# model space, with 16 initial filters and 8 stacked cells.
# The network is specialized for CIFAR-10 dataset with 32x32 input resolution.
#
# The :class:`~nni.retiarii.hub.pytorch.DARTS` model space here is provided by :doc:`model space hub </nas/space_hub>`,
# where we have supported multiple popular model spaces for plug-and-play.
#
# .. tip::
#
# The model space here can be replaced with any space provided in the hub,
# or even customized spaces built from scratch.
model_space = DartsSpace(16, 8, 'cifar')
# %%
#
# Search on the model space
# -------------------------
#
# .. warning::
#
# Please set ``fast_dev_run`` to False to reproduce the our claimed results.
# Otherwise, only a few mini-batches will be run.
fast_dev_run = True
# %%
#
# Evaluator
# ^^^^^^^^^
#
# To begin exploring the model space, one firstly need to have an evaluator to provide the criterion of a "good model".
# As we are searching on CIFAR-10 dataset, one can easily use the :class:`~nni.retiarii.evaluator.pytorch.Classification`
# as a starting point.
#
# Note that for a typical setup of NAS, the model search should be on validation set, and the evaluation of the final searched model
# should be on test set. However, as CIFAR-10 dataset doesn't have a test dataset (only 50k train + 10k valid),
# we have to split the original training set into a training set and a validation set.
# The recommended train/val split by `DARTS`_ strategy is 1:1.
import numpy as np
from nni.retiarii.evaluator.pytorch import Classification
from torch.utils.data import SubsetRandomSampler
transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
train_data = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform)
num_samples = len(train_data)
indices = np.random.permutation(num_samples)
split = num_samples // 2
search_train_loader = DataLoader(
train_data, batch_size=64, num_workers=6,
sampler=SubsetRandomSampler(indices[:split]),
)
search_valid_loader = DataLoader(
train_data, batch_size=64, num_workers=6,
sampler=SubsetRandomSampler(indices[split:]),
)
evaluator = Classification(
learning_rate=1e-3,
weight_decay=1e-4,
train_dataloaders=search_train_loader,
val_dataloaders=search_valid_loader,
max_epochs=10,
gpus=1,
fast_dev_run=fast_dev_run,
)
# %%
#
# Strategy
# ^^^^^^^^
#
# We will use `DARTS`_ (Differentiable ARchiTecture Search) as the search strategy to explore the model space.
# :class:`~nni.retiarii.strategy.DARTS` strategy belongs to the category of :ref:`one-shot strategy <one-shot-nas>`.
# The fundamental differences between One-shot strategies and :ref:`multi-trial strategies <multi-trial-nas>` is that,
# one-shot strategy combines search with model training into a single run.
# Compared to multi-trial strategies, one-shot NAS doesn't need to iteratively spawn new trials (i.e., models),
# and thus saves the excessive cost of model training.
# It's worth mentioning that one-shot NAS also suffers from multiple drawbacks despite its computational efficiency.
# We recommend
# `Weight-Sharing Neural Architecture Search: A Battle to Shrink the Optimization Gap <https://arxiv.org/abs/2008.01475>`__
# and
# `How Does Supernet Help in Neural Architecture Search? <https://arxiv.org/abs/2010.08219>`__ for interested readers.
#
# :class:`~nni.retiarii.strategy.DARTS` strategy is provided as one of NNI's :doc:`built-in search strategies </nas/exploration_strategy>`.
# Using it can be as simple as one line of code.
from nni.retiarii.strategy import DARTS as DartsStrategy
strategy = DartsStrategy()
# %%
#
# .. tip:: The ``DartsStrategy`` here can be replaced by any search strategies, even multi-trial strategies.
#
# If you want to know how DARTS strategy works, here is a brief version.
# Under the hood, DARTS converts the cell into a densely connected graph, and put operators on edges (see the following figure).
# Since the operators are not decided yet, every edge is a weighted mixture of multiple operators (multiple color in the figure).
# DARTS then learns to assign the optimal "color" for each edge during the network training.
# It finally selects one "color" for each edge, and drops redundant edges.
# The weights on the edges are called *architecture weights*.
#
# .. image:: ../../img/darts_illustration.png
#
# It's NOT reflected in the figure that, for DARTS model space, exactly two inputs are kept for every node.
#
# Launch experiment
# ^^^^^^^^^^^^^^^^^
#
# We then come to the step of launching the experiment.
# This step is similar to what we have done in the :doc:`beginner tutorial <hello_nas>`,
# except that the ``execution_engine`` argument should be set to ``oneshot``.
from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
config = RetiariiExeConfig(execution_engine='oneshot')
experiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)
experiment.run(config)
# %%
#
# .. tip::
#
# The search process can be visualized with tensorboard. For example::
#
# tensorboard --logdir=./lightning_logs
#
# Then, open the browser and go to http://localhost:6006/ to monitor the search process.
#
# .. image:: ../../img/darts_search_process.png
#
# We can then retrieve the best model found by the strategy with ``export_top_models``.
# Here, the retrieved model is a dict (called *architecture dict*) describing the selected normal cell and reduction cell.
exported_arch = experiment.export_top_models()[0]
exported_arch
# %%
#
# The cell can be visualized with the following code snippet
# (copied and modified from `DARTS visualization <https://github.com/quark0/darts/blob/master/cnn/visualize.py>`__).
import io
import graphviz
import matplotlib.pyplot as plt
from PIL import Image
def plot_single_cell(arch_dict, cell_name):
g = graphviz.Digraph(
node_attr=dict(style='filled', shape='rect', align='center'),
format='png'
)
g.body.extend(['rankdir=LR'])
g.node('c_{k-2}', fillcolor='darkseagreen2')
g.node('c_{k-1}', fillcolor='darkseagreen2')
assert len(arch_dict) % 2 == 0
for i in range(2, 6):
g.node(str(i), fillcolor='lightblue')
for i in range(2, 6):
for j in range(2):
op = arch_dict[f'{cell_name}/op_{i}_{j}']
from_ = arch_dict[f'{cell_name}/input_{i}_{j}']
if from_ == 0:
u = 'c_{k-2}'
elif from_ == 1:
u = 'c_{k-1}'
else:
u = str(from_)
v = str(i)
g.edge(u, v, label=op, fillcolor='gray')
g.node('c_{k}', fillcolor='palegoldenrod')
for i in range(2, 6):
g.edge(str(i), 'c_{k}', fillcolor='gray')
g.attr(label=f'{cell_name.capitalize()} cell')
image = Image.open(io.BytesIO(g.pipe()))
return image
def plot_double_cells(arch_dict):
image1 = plot_single_cell(arch_dict, 'normal')
image2 = plot_single_cell(arch_dict, 'reduce')
height_ratio = max(image1.size[1] / image1.size[0], image2.size[1] / image2.size[0])
_, axs = plt.subplots(1, 2, figsize=(20, 10 * height_ratio))
axs[0].imshow(image1)
axs[1].imshow(image2)
axs[0].axis('off')
axs[1].axis('off')
plt.show()
plot_double_cells(exported_arch)
# %%
#
# .. warning:: The cell above is obtained via ``fast_dev_run`` (i.e., running only 1 mini-batch).
#
# When ``fast_dev_run`` is turned off, we get a model with the following architecture,
# where you might notice an interesting fact that around half the operations have selected ``sep_conv_3x3``.
plot_double_cells({
'normal/op_2_0': 'sep_conv_3x3',
'normal/input_2_0': 1,
'normal/op_2_1': 'sep_conv_3x3',
'normal/input_2_1': 0,
'normal/op_3_0': 'sep_conv_3x3',
'normal/input_3_0': 1,
'normal/op_3_1': 'sep_conv_3x3',
'normal/input_3_1': 2,
'normal/op_4_0': 'sep_conv_3x3',
'normal/input_4_0': 1,
'normal/op_4_1': 'sep_conv_3x3',
'normal/input_4_1': 0,
'normal/op_5_0': 'sep_conv_3x3',
'normal/input_5_0': 1,
'normal/op_5_1': 'max_pool_3x3',
'normal/input_5_1': 0,
'reduce/op_2_0': 'sep_conv_3x3',
'reduce/input_2_0': 0,
'reduce/op_2_1': 'sep_conv_3x3',
'reduce/input_2_1': 1,
'reduce/op_3_0': 'dil_conv_5x5',
'reduce/input_3_0': 2,
'reduce/op_3_1': 'sep_conv_3x3',
'reduce/input_3_1': 0,
'reduce/op_4_0': 'dil_conv_5x5',
'reduce/input_4_0': 2,
'reduce/op_4_1': 'sep_conv_5x5',
'reduce/input_4_1': 1,
'reduce/op_5_0': 'sep_conv_5x5',
'reduce/input_5_0': 4,
'reduce/op_5_1': 'dil_conv_5x5',
'reduce/input_5_1': 2
})
# %%
#
# Retrain the searched model
# --------------------------
#
# What we have got in the last step, is only a cell structure.
# To get a final usable model with trained weights, we need to construct a real model based on this structure,
# and then fully train it.
#
# To construct a fixed model based on the architecture dict exported from the experiment,
# we can use :func:`nni.retiarii.fixed_arch`. Seemingly, we are still creating a space.
# But under the with-context, we are actually creating a fixed model.
from nni.retiarii import fixed_arch
with fixed_arch(exported_arch):
final_model = DartsSpace(16, 8, 'cifar')
# %%
#
# We then train the model on full CIFAR-10 training dataset, and evaluate it on the original CIFAR-10 validation dataset.
train_loader = DataLoader(train_data, batch_size=96, num_workers=6) # Use the original training data
# %%
#
# The validation data loader can be reused.
valid_loader
# %%
#
# We must create a new evaluator here because a different data split is used.
# Also, we should avoid the underlying pytorch-lightning implementation of :class:`~nni.retiarii.evaluator.pytorch.Classification`
# evaluator from loading the wrong checkpoint.
max_epochs = 100
evaluator = Classification(
learning_rate=1e-3,
weight_decay=1e-4,
train_dataloaders=train_loader,
val_dataloaders=valid_loader,
max_epochs=max_epochs,
gpus=1,
export_onnx=False, # Disable ONNX export for this experiment
fast_dev_run=fast_dev_run, # Should be false for fully training
)
evaluator.fit(final_model)
# %%
#
# .. note:: When ``fast_dev_run`` is turned off, we achieve a validation accuracy of 89.69% after training for 100 epochs.
#
# Reproduce results in DARTS paper
# --------------------------------
#
# After a brief walkthrough of search + retrain process with one-shot strategy,
# we then fill the gap between our results (89.69%) and the results in the `DARTS` paper.
# This is because we didn't introduce some extra training tricks, including `DropPath <https://arxiv.org/pdf/1605.07648v4.pdf>`__,
# Auxiliary loss, gradient clipping and augmentations like `Cutout <https://arxiv.org/pdf/1708.04552v2.pdf>`__.
# They also train the deeper (20 cells) and wider (36 channels) networks for longer time (600 epochs).
#
#
# Evaluator
# ^^^^^^^^^
#
# To implement these tricks, we first need to rewrite a few parts of evaluator.
#
# Working with one-shot strategies, evaluators need to be implemented in the style of :ref:`PyTorch-Lightning <lightning-evaluator>`,
# The full tutorial can be found in :doc:`/nas/evaluator`.
# Putting it briefly, the core part of writing a new evaluator is to write a new LightningModule.
# `LightingModule <https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html>`__ is a concept in
# PyTorch-Lightning, which organizes the model training process into a list of functions, such as,
# ``training_step``, ``validation_step``, ``configure_optimizers``, etc.
# Since we are merely adding a few ingredients to :class:`~nni.retiarii.evaluator.pytorch.Classification`,
# we can simply inherit :class:`~nni.retiarii.evaluator.pytorch.ClassificationModule`, which is the underlying LightningModule
# behind :class:`~nni.retiarii.evaluator.pytorch.Classification`.
# This could look intimidating at first, but most of them are just plug-and-play tricks which you don't need to know details about.
import torch
from nni.retiarii.evaluator.pytorch import ClassificationModule
class DartsClassificationModule(ClassificationModule):
def __init__(
self,
learning_rate: float = 0.001,
weight_decay: float = 0.,
auxiliary_loss_weight: float = 0.4,
max_epochs: int = 600
):
self.auxiliary_loss_weight = auxiliary_loss_weight
# Training length will be used in LR scheduler
self.max_epochs = max_epochs
super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False)
def configure_optimizers(self):
"""Customized optimizer with momentum, as well as a scheduler."""
optimizer = torch.optim.SGD(
self.parameters(),
momentum=0.9,
lr=self.hparams.learning_rate,
weight_decay=self.hparams.weight_decay
)
return {
'optimizer': optimizer,
'lr_scheduler': torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, self.max_epochs, eta_min=1e-3)
}
def training_step(self, batch, batch_idx):
"""Training step, customized with auxiliary loss."""
x, y = batch
if self.auxiliary_loss_weight:
y_hat, y_aux = self(x)
loss_main = self.criterion(y_hat, y)
loss_aux = self.criterion(y_aux, y)
self.log('train_loss_main', loss_main)
self.log('train_loss_aux', loss_aux)
loss = loss_main + self.auxiliary_loss_weight * loss_aux
else:
y_hat = self(x)
loss = self.criterion(y_hat, y)
self.log('train_loss', loss, prog_bar=True)
for name, metric in self.metrics.items():
self.log('train_' + name, metric(y_hat, y), prog_bar=True)
return loss
def on_train_epoch_start(self):
# Set drop path probability before every epoch. This has no effect if drop path is not enabled in model.
self.model.set_drop_path_prob(self.model.drop_path_prob * self.current_epoch / self.max_epochs)
# Logging learning rate at the beginning of every epoch
self.log('lr', self.trainer.optimizers[0].param_groups[0]['lr'])
# %%
#
# The full evaluator is written as follows,
# which simply wraps everything (except model space and search strategy of course), in a single object.
# :class:`~nni.retiarii.evaluator.pytorch.Lightning` here is a special type of evaluator.
# Don't forget to use the train/val data split specialized for search (1:1) here.
from nni.retiarii.evaluator.pytorch import Lightning, Trainer
max_epochs = 50
evaluator = Lightning(
DartsClassificationModule(0.025, 3e-4, 0., max_epochs),
Trainer(
gpus=1,
max_epochs=max_epochs,
fast_dev_run=fast_dev_run,
),
train_dataloaders=search_train_loader,
val_dataloaders=search_valid_loader
)
# %%
#
# Strategy
# ^^^^^^^^
#
# :class:`~nni.retiarii.strategy.DARTS` strategy is created with gradient clip turned on.
# If you are familiar with PyTorch-Lightning, you might aware that gradient clipping can be enabled in Lightning trainer.
# However, enabling gradient cip in the trainer above won't work, because the underlying
# implementation of :class:`~nni.retiarii.strategy.DARTS` strategy is based on
# `manual optimization <https://pytorch-lightning.readthedocs.io/en/stable/common/optimization.html>`__.
strategy = DartsStrategy(gradient_clip_val=5.)
# %%
#
# Launch experiment
# ^^^^^^^^^^^^^^^^^
#
# Then we use the newly created evaluator and strategy to launch the experiment again.
#
# .. warning::
#
# ``model_space`` has to be re-instantiated because a known limitation,
# i.e., one model space can't be reused across multiple experiments.
model_space = DartsSpace(16, 8, 'cifar')
config = RetiariiExeConfig(execution_engine='oneshot')
experiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)
experiment.run(config)
exported_arch = experiment.export_top_models()[0]
exported_arch
# %%
#
# We get the following architecture when ``fast_dev_run`` is set to False. It takes around 8 hours on a P100 GPU.
plot_double_cells({
'normal/op_2_0': 'sep_conv_3x3',
'normal/input_2_0': 0,
'normal/op_2_1': 'sep_conv_3x3',
'normal/input_2_1': 1,
'normal/op_3_0': 'sep_conv_3x3',
'normal/input_3_0': 1,
'normal/op_3_1': 'skip_connect',
'normal/input_3_1': 0,
'normal/op_4_0': 'sep_conv_3x3',
'normal/input_4_0': 0,
'normal/op_4_1': 'max_pool_3x3',
'normal/input_4_1': 1,
'normal/op_5_0': 'sep_conv_3x3',
'normal/input_5_0': 0,
'normal/op_5_1': 'sep_conv_3x3',
'normal/input_5_1': 1,
'reduce/op_2_0': 'max_pool_3x3',
'reduce/input_2_0': 0,
'reduce/op_2_1': 'sep_conv_5x5',
'reduce/input_2_1': 1,
'reduce/op_3_0': 'dil_conv_5x5',
'reduce/input_3_0': 2,
'reduce/op_3_1': 'max_pool_3x3',
'reduce/input_3_1': 0,
'reduce/op_4_0': 'max_pool_3x3',
'reduce/input_4_0': 0,
'reduce/op_4_1': 'sep_conv_3x3',
'reduce/input_4_1': 2,
'reduce/op_5_0': 'max_pool_3x3',
'reduce/input_5_0': 0,
'reduce/op_5_1': 'skip_connect',
'reduce/input_5_1': 2
})
# %%
#
# Retrain
# ^^^^^^^
#
# When retraining,
# we extend the original dataloader to introduce another trick called `Cutout <https://arxiv.org/pdf/1708.04552v2.pdf>`__.
# Cutout is a data augmentation technique that randomly masks out rectangular regions in images.
# In CIFAR-10, the typical masked size is 16x16 (the image sizes are 32x32 in the dataset).
def cutout_transform(img, length: int = 16):
h, w = img.size(1), img.size(2)
mask = np.ones((h, w), np.float32)
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - length // 2, 0, h)
y2 = np.clip(y + length // 2, 0, h)
x1 = np.clip(x - length // 2, 0, w)
x2 = np.clip(x + length // 2, 0, w)
mask[y1: y2, x1: x2] = 0.
mask = torch.from_numpy(mask)
mask = mask.expand_as(img)
img *= mask
return img
transform_with_cutout = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
cutout_transform,
])
# %%
#
# The train dataloader needs to be reinstantiated with the new transform.
# The validation dataloader is not affected, and thus can be reused.
train_data_cutout = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform_with_cutout)
train_loader_cutout = DataLoader(train_data_cutout, batch_size=96)
# %%
#
# We then create the final model based on the new exported architecture.
# This time, auxiliary loss and drop path probability is enabled.
#
# Following the same procedure as paper, we also increase the number of filters to 36, and number of cells to 20,
# so as to reasonably increase the model size and boost the performance.
with fixed_arch(exported_arch):
final_model = DartsSpace(36, 20, 'cifar', auxiliary_loss=True, drop_path_prob=0.2)
# %%
#
# We create a new evaluator for the retraining process, where the gradient clipping is put into the keyword arguments of trainer.
max_epochs = 600
evaluator = Lightning(
DartsClassificationModule(0.025, 3e-4, 0.4, max_epochs),
Trainer(
gpus=1,
gradient_clip_val=5.,
max_epochs=max_epochs,
fast_dev_run=fast_dev_run
),
train_dataloaders=train_loader_cutout,
val_dataloaders=valid_loader,
)
evaluator.fit(final_model)
# %%
#
# When ``fast_dev_run`` is turned off, after retraining, the architecture yields a top-1 accuracy of 97.12%.
# If we take the best snapshot throughout the retrain process,
# there is a chance that the top-1 accuracy will be 97.28%.
#
# .. image:: ../../img/darts_val_acc.png
#
# In the figure, the orange line is the validation accuracy curve after training for 600 epochs,
# while the red line corresponding the previous version in this tutorial before adding all the training tricks and
# only trains for 100 epochs.
#
# The results outperforms "DARTS (first order) + cutout" in `DARTS`_ paper, which is only 97.00±0.14%.
# It's even comparable with "DARTS (second order) + cutout" in the paper (97.24±0.09%),
# though we didn't implement the second order version.
# The implementation of second order DARTS is in our future plan, and we also welcome your contribution.
docs/source/tutorials/darts.py.md5 0 → 100644
View file @ bddaf36d
240d9ba3c97be549376aa4ef2bd08344
\ No newline at end of file
docs/source/tutorials/darts.rst 0 → 100644
View file @ bddaf36d
.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "tutorials/darts.py"
.. LINE NUMBERS ARE GIVEN BELOW.
.. only:: html
.. note::
:class: sphx-glr-download-link-note
Click :ref:`here <sphx_glr_download_tutorials_darts.py>`
to download the full example code
.. rst-class:: sphx-glr-example-title
.. _sphx_glr_tutorials_darts.py:
Searching in DARTS search space
===============================
In this tutorial, we demonstrate how to search in the famous model space proposed in `DARTS`_.
Through this process, you will learn:
* How to use the built-in model spaces from NNI's model space hub.
* How to use one-shot exploration strategies to explore a model space.
* How to customize evaluators to achieve the best performance.
In the end, we get a strong-performing model on CIFAR-10 dataset, which achieves up to 97.28% accuracy.
.. attention::
Running this tutorial requires a GPU.
If you don't have one, you can set ``gpus`` in :class:`~nni.retiarii.evaluator.pytorch.Classification` to be 0,
but do note that it will be much slower.
.. _DARTS: https://arxiv.org/abs/1806.09055
Use a pre-searched model
------------------------
Similar to `the beginner tutorial of PyTorch <https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html>`__,
we begin with CIFAR-10 dataset, which is a image classification dataset of 10 categories.
The images in CIFAR-10 are of size 3x32x32, i.e., RGB-colored images of 32x32 pixels in size.
We first load the CIFAR-10 dataset with torchvision.
.. GENERATED FROM PYTHON SOURCE LINES 32-49
.. code-block:: default
import nni
import torch
from torchvision import transforms
from torchvision.datasets import CIFAR10
from nni.retiarii.evaluator.pytorch import DataLoader
CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]
transform_valid = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
valid_data = nni.trace(CIFAR10)(root='./data', train=False, download=True, transform=transform_valid)
valid_loader = DataLoader(valid_data, batch_size=256, num_workers=6)
.. rst-class:: sphx-glr-script-out
.. code-block:: none
Files already downloaded and verified
.. GENERATED FROM PYTHON SOURCE LINES 50-74
.. note::
If you are to use multi-trial strategies, wrapping CIFAR10 with :func:`nni.trace` and
use DataLoader from ``nni.retiarii.evaluator.pytorch`` (instead of ``torch.utils.data``) are mandatory.
Otherwise, it's optional.
When working with famous datasets like CIFAR-10 or ImageNet,
it's tempting to use or finetune from a pretrained model, like ResNet.
There's nothing wrong with doing so, and sometimes it might be beneficial.
Thanks to the development of NAS, we now have quite a large number of *pre-searched models*,
which are produced by most popular NAS literatures.
You can easily load these models, validate their performances, and finetune them if you need.
We present :doc:`model space hub </nas/space_hub>`, where you can find many built-in model spaces,
along with many pre-searched models.
We choose one from `DARTS`_ search space, which is natively trained on our target dataset, CIFAR-10,
so as to save the tedious steps of finetuning.
.. tip::
Finetuning a pre-searched model on other datasets is no different from finetuning *any model*.
We recommend reading
`this tutorial of object detection finetuning <https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html>`__
if you want to know how finetuning is generally done in PyTorch.
.. GENERATED FROM PYTHON SOURCE LINES 75-97
.. code-block:: default
from nni.retiarii.hub.pytorch import DARTS as DartsSpace
darts_v2_model = DartsSpace.load_searched_model('darts-v2', pretrained=True, download=True)
def evaluate_model(model, cuda=False):
device = torch.device('cuda' if cuda else 'cpu')
model.to(device)
model.eval()
with torch.no_grad():
correct = total = 0
for inputs, targets in valid_loader:
inputs, targets = inputs.to(device), targets.to(device)
logits = model(inputs)
_, predict = torch.max(logits, 1)
correct += (predict == targets).sum().cpu().item()
total += targets.size(0)
print('Accuracy:', correct / total)
return correct / total
evaluate_model(darts_v2_model, True) # Set this to false if there's no GPU.
.. rst-class:: sphx-glr-script-out
.. code-block:: none
Accuracy: 0.9737
0.9737
.. GENERATED FROM PYTHON SOURCE LINES 98-162
The journey could end here. Or you are interested,
we can go a step further to search a model within :class:`~nni.retiarii.hub.pytorch.DARTS` space on our own.
Use the model space
-------------------
The model space provided in `DARTS`_ originated from `NASNet <https://arxiv.org/abs/1707.07012>`__,
where the full model is constructed by repeatedly stacking a single computational unit (called a **cell**).
There are two types of cells within a network. The first type is called *normal cell*, and the second type is called *reduction cell*.
The key difference between normal and reduction cell is that the reduction cell will downsample the input feature map,
and decrease its resolution. Normal and reduction cells are stacked alternately, as shown in the following figure.
.. image:: ../../img/nasnet_cell_stack.png
A cell takes outputs from two previous cells as inputs and contains a collection of *nodes*.
Each node takes two previous nodes within the same cell (or the two cell inputs),
and applies an *operator* (e.g., convolution, or max-pooling) to each input,
and sums the outputs of operators as the output of the node.
The output of cell is the concatenation of all the nodes that are never used as inputs of another node.
We recommend reading `NDS <https://arxiv.org/pdf/1905.13214.pdf>`__ or `ENAS <https://arxiv.org/abs/1802.03268>`__ for details.
We illustrate an example of cells in the following figure.
.. image:: ../../img/nasnet_cell.png
The search space proposed in `DARTS`_ paper introduced two modifications to the original space
in `NASNet <https://arxiv.org/abs/1707.07012>`__.
Firstly, the operator candidates have been narrowed down to seven:
- Max pooling 3x3
- Average pooling 3x3
- Skip connect (Identity)
- Separable convolution 3x3
- Separable convolution 5x5
- Dilated convolution 3x3
- Dilated convolution 5x5
Secondly, the output of cell is the concatenate of **all the nodes within the cell**.
As the search space is based on cell, once the normal and reduction cell has been fixed, we can stack them for indefinite times.
To save the search cost, the common practice is to reduce the number of filters (i.e., channels) and number of stacked cells
during the search phase, and increase them back when training the final searched architecture.
.. note::
`DARTS`_ is one of those papers that innovate both in search space and search strategy.
In this tutorial, we will search on **model space** provided by DARTS with **search strategy** proposed by DARTS.
We refer to them as *DARTS model space* (``DartsSpace``) and *DARTS strategy* (``DartsStrategy``), respectively.
We did NOT imply that the :class:`~nni.retiarii.hub.pytorch.DARTS` space and
:class:`~nni.retiarii.strategy.DARTS` strategy has to used together.
You can always explore the DARTS space with another search strategy, or use your own strategy to search a different model space.
In the following example, we initialize a :class:`~nni.retiarii.hub.pytorch.DARTS`
model space, with 16 initial filters and 8 stacked cells.
The network is specialized for CIFAR-10 dataset with 32x32 input resolution.
The :class:`~nni.retiarii.hub.pytorch.DARTS` model space here is provided by :doc:`model space hub </nas/space_hub>`,
where we have supported multiple popular model spaces for plug-and-play.
.. tip::
The model space here can be replaced with any space provided in the hub,
or even customized spaces built from scratch.
.. GENERATED FROM PYTHON SOURCE LINES 163-166
.. code-block:: default
model_space = DartsSpace(16, 8, 'cifar')
.. GENERATED FROM PYTHON SOURCE LINES 167-174
Search on the model space
-------------------------
.. warning::
Please set ``fast_dev_run`` to False to reproduce the our claimed results.
Otherwise, only a few mini-batches will be run.
.. GENERATED FROM PYTHON SOURCE LINES 175-178
.. code-block:: default
fast_dev_run = True
.. GENERATED FROM PYTHON SOURCE LINES 179-190
Evaluator
^^^^^^^^^
To begin exploring the model space, one firstly need to have an evaluator to provide the criterion of a "good model".
As we are searching on CIFAR-10 dataset, one can easily use the :class:`~nni.retiarii.evaluator.pytorch.Classification`
as a starting point.
Note that for a typical setup of NAS, the model search should be on validation set, and the evaluation of the final searched model
should be on test set. However, as CIFAR-10 dataset doesn't have a test dataset (only 50k train + 10k valid),
we have to split the original training set into a training set and a validation set.
The recommended train/val split by `DARTS`_ strategy is 1:1.
.. GENERATED FROM PYTHON SOURCE LINES 191-229
.. code-block:: default
import numpy as np
from nni.retiarii.evaluator.pytorch import Classification
from torch.utils.data import SubsetRandomSampler
transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
train_data = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform)
num_samples = len(train_data)
indices = np.random.permutation(num_samples)
split = num_samples // 2
search_train_loader = DataLoader(
train_data, batch_size=64, num_workers=6,
sampler=SubsetRandomSampler(indices[:split]),
)
search_valid_loader = DataLoader(
train_data, batch_size=64, num_workers=6,
sampler=SubsetRandomSampler(indices[split:]),
)
evaluator = Classification(
learning_rate=1e-3,
weight_decay=1e-4,
train_dataloaders=search_train_loader,
val_dataloaders=search_valid_loader,
max_epochs=10,
gpus=1,
fast_dev_run=fast_dev_run,
)
.. rst-class:: sphx-glr-script-out
.. code-block:: none
Files already downloaded and verified
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:445: LightningDeprecationWarning: Setting `Trainer(gpus=1)` is deprecated in v1.7 and will be removed in v2.0. Please use `Trainer(accelerator='gpu', devices=1)` instead.
rank_zero_deprecation(
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
Running in `fast_dev_run` mode: will run the requested loop using 1 batch(es). Logging and checkpointing is suppressed.
`Trainer(limit_train_batches=1)` was configured so 1 batch per epoch will be used.
`Trainer(limit_val_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_test_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_predict_batches=1)` was configured so 1 batch will be used.
`Trainer(val_check_interval=1.0)` was configured so validation will run at the end of the training epoch..
.. GENERATED FROM PYTHON SOURCE LINES 230-247
Strategy
^^^^^^^^
We will use `DARTS`_ (Differentiable ARchiTecture Search) as the search strategy to explore the model space.
:class:`~nni.retiarii.strategy.DARTS` strategy belongs to the category of :ref:`one-shot strategy <one-shot-nas>`.
The fundamental differences between One-shot strategies and :ref:`multi-trial strategies <multi-trial-nas>` is that,
one-shot strategy combines search with model training into a single run.
Compared to multi-trial strategies, one-shot NAS doesn't need to iteratively spawn new trials (i.e., models),
and thus saves the excessive cost of model training.
It's worth mentioning that one-shot NAS also suffers from multiple drawbacks despite its computational efficiency.
We recommend
`Weight-Sharing Neural Architecture Search: A Battle to Shrink the Optimization Gap <https://arxiv.org/abs/2008.01475>`__
and
`How Does Supernet Help in Neural Architecture Search? <https://arxiv.org/abs/2010.08219>`__ for interested readers.
:class:`~nni.retiarii.strategy.DARTS` strategy is provided as one of NNI's :doc:`built-in search strategies </nas/exploration_strategy>`.
Using it can be as simple as one line of code.
.. GENERATED FROM PYTHON SOURCE LINES 248-253
.. code-block:: default
from nni.retiarii.strategy import DARTS as DartsStrategy
strategy = DartsStrategy()
.. GENERATED FROM PYTHON SOURCE LINES 254-273
.. tip:: The ``DartsStrategy`` here can be replaced by any search strategies, even multi-trial strategies.
If you want to know how DARTS strategy works, here is a brief version.
Under the hood, DARTS converts the cell into a densely connected graph, and put operators on edges (see the following figure).
Since the operators are not decided yet, every edge is a weighted mixture of multiple operators (multiple color in the figure).
DARTS then learns to assign the optimal "color" for each edge during the network training.
It finally selects one "color" for each edge, and drops redundant edges.
The weights on the edges are called *architecture weights*.
.. image:: ../../img/darts_illustration.png
It's NOT reflected in the figure that, for DARTS model space, exactly two inputs are kept for every node.
Launch experiment
^^^^^^^^^^^^^^^^^
We then come to the step of launching the experiment.
This step is similar to what we have done in the :doc:`beginner tutorial <hello_nas>`,
except that the ``execution_engine`` argument should be set to ``oneshot``.
.. GENERATED FROM PYTHON SOURCE LINES 274-281
.. code-block:: default
from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
config = RetiariiExeConfig(execution_engine='oneshot')
experiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)
experiment.run(config)
.. rst-class:: sphx-glr-script-out
.. code-block:: none
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
| Name | Type | Params
-----------------------------------------------
0 | model | ClassificationModule | 3.0 M
-----------------------------------------------
3.0 M Trainable params
0 Non-trainable params
3.0 M Total params
12.164 Total estimated model params size (MB)
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1891: PossibleUserWarning: The number of training batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
rank_zero_warn(
Training: 0it [00:00, ?it/s] Training: 0%| | 0/1 [00:00<?, ?it/s] Epoch 0: 0%| | 0/1 [00:00<?, ?it/s] Epoch 0: 100%|##########| 1/1 [00:03<00:00, 3.24s/it] Epoch 0: 100%|##########| 1/1 [00:03<00:00, 3.25s/it, v_num=, train_loss=2.410, train_acc=0.141] Epoch 0: 100%|##########| 1/1 [00:03<00:00, 3.25s/it, v_num=, train_loss=2.410, train_acc=0.141]`Trainer.fit` stopped: `max_epochs=1` reached.
Epoch 0: 100%|##########| 1/1 [00:03<00:00, 3.26s/it, v_num=, train_loss=2.410, train_acc=0.141]
.. GENERATED FROM PYTHON SOURCE LINES 282-294
.. tip::
The search process can be visualized with tensorboard. For example::
tensorboard --logdir=./lightning_logs
Then, open the browser and go to http://localhost:6006/ to monitor the search process.
.. image:: ../../img/darts_search_process.png
We can then retrieve the best model found by the strategy with ``export_top_models``.
Here, the retrieved model is a dict (called *architecture dict*) describing the selected normal cell and reduction cell.
.. GENERATED FROM PYTHON SOURCE LINES 295-300
.. code-block:: default
exported_arch = experiment.export_top_models()[0]
exported_arch
.. rst-class:: sphx-glr-script-out
.. code-block:: none
{'normal/op_2_0': 'avg_pool_3x3', 'normal/input_2_0': 1, 'normal/op_2_1': 'sep_conv_5x5', 'normal/input_2_1': 0, 'normal/op_3_0': 'sep_conv_3x3', 'normal/input_3_0': 1, 'normal/op_3_1': 'dil_conv_5x5', 'normal/input_3_1': 2, 'normal/op_4_0': 'sep_conv_3x3', 'normal/input_4_0': 0, 'normal/op_4_1': 'dil_conv_5x5', 'normal/input_4_1': 3, 'normal/op_5_0': 'dil_conv_3x3', 'normal/input_5_0': 0, 'normal/op_5_1': 'dil_conv_5x5', 'normal/input_5_1': 4, 'reduce/op_2_0': 'max_pool_3x3', 'reduce/input_2_0': 1, 'reduce/op_2_1': 'sep_conv_5x5', 'reduce/input_2_1': 0, 'reduce/op_3_0': 'skip_connect', 'reduce/input_3_0': 1, 'reduce/op_3_1': 'dil_conv_3x3', 'reduce/input_3_1': 2, 'reduce/op_4_0': 'sep_conv_5x5', 'reduce/input_4_0': 0, 'reduce/op_4_1': 'dil_conv_3x3', 'reduce/input_4_1': 1, 'reduce/op_5_0': 'avg_pool_3x3', 'reduce/input_5_0': 1, 'reduce/op_5_1': 'dil_conv_3x3', 'reduce/input_5_1': 4}
.. GENERATED FROM PYTHON SOURCE LINES 301-303
The cell can be visualized with the following code snippet
(copied and modified from `DARTS visualization <https://github.com/quark0/darts/blob/master/cnn/visualize.py>`__).
.. GENERATED FROM PYTHON SOURCE LINES 304-359
.. code-block:: default
import io
import graphviz
import matplotlib.pyplot as plt
from PIL import Image
def plot_single_cell(arch_dict, cell_name):
g = graphviz.Digraph(
node_attr=dict(style='filled', shape='rect', align='center'),
format='png'
)
g.body.extend(['rankdir=LR'])
g.node('c_{k-2}', fillcolor='darkseagreen2')
g.node('c_{k-1}', fillcolor='darkseagreen2')
assert len(arch_dict) % 2 == 0
for i in range(2, 6):
g.node(str(i), fillcolor='lightblue')
for i in range(2, 6):
for j in range(2):
op = arch_dict[f'{cell_name}/op_{i}_{j}']
from_ = arch_dict[f'{cell_name}/input_{i}_{j}']
if from_ == 0:
u = 'c_{k-2}'
elif from_ == 1:
u = 'c_{k-1}'
else:
u = str(from_)
v = str(i)
g.edge(u, v, label=op, fillcolor='gray')
g.node('c_{k}', fillcolor='palegoldenrod')
for i in range(2, 6):
g.edge(str(i), 'c_{k}', fillcolor='gray')
g.attr(label=f'{cell_name.capitalize()} cell')
image = Image.open(io.BytesIO(g.pipe()))
return image
def plot_double_cells(arch_dict):
image1 = plot_single_cell(arch_dict, 'normal')
image2 = plot_single_cell(arch_dict, 'reduce')
height_ratio = max(image1.size[1] / image1.size[0], image2.size[1] / image2.size[0])
_, axs = plt.subplots(1, 2, figsize=(20, 10 * height_ratio))
axs[0].imshow(image1)
axs[1].imshow(image2)
axs[0].axis('off')
axs[1].axis('off')
plt.show()
plot_double_cells(exported_arch)
.. image-sg:: /tutorials/images/sphx_glr_darts_001.png
:alt: darts
:srcset: /tutorials/images/sphx_glr_darts_001.png
:class: sphx-glr-single-img
.. GENERATED FROM PYTHON SOURCE LINES 360-364
.. warning:: The cell above is obtained via ``fast_dev_run`` (i.e., running only 1 mini-batch).
When ``fast_dev_run`` is turned off, we get a model with the following architecture,
where you might notice an interesting fact that around half the operations have selected ``sep_conv_3x3``.
.. GENERATED FROM PYTHON SOURCE LINES 365-401
.. code-block:: default
plot_double_cells({
'normal/op_2_0': 'sep_conv_3x3',
'normal/input_2_0': 1,
'normal/op_2_1': 'sep_conv_3x3',
'normal/input_2_1': 0,
'normal/op_3_0': 'sep_conv_3x3',
'normal/input_3_0': 1,
'normal/op_3_1': 'sep_conv_3x3',
'normal/input_3_1': 2,
'normal/op_4_0': 'sep_conv_3x3',
'normal/input_4_0': 1,
'normal/op_4_1': 'sep_conv_3x3',
'normal/input_4_1': 0,
'normal/op_5_0': 'sep_conv_3x3',
'normal/input_5_0': 1,
'normal/op_5_1': 'max_pool_3x3',
'normal/input_5_1': 0,
'reduce/op_2_0': 'sep_conv_3x3',
'reduce/input_2_0': 0,
'reduce/op_2_1': 'sep_conv_3x3',
'reduce/input_2_1': 1,
'reduce/op_3_0': 'dil_conv_5x5',
'reduce/input_3_0': 2,
'reduce/op_3_1': 'sep_conv_3x3',
'reduce/input_3_1': 0,
'reduce/op_4_0': 'dil_conv_5x5',
'reduce/input_4_0': 2,
'reduce/op_4_1': 'sep_conv_5x5',
'reduce/input_4_1': 1,
'reduce/op_5_0': 'sep_conv_5x5',
'reduce/input_5_0': 4,
'reduce/op_5_1': 'dil_conv_5x5',
'reduce/input_5_1': 2
})
.. image-sg:: /tutorials/images/sphx_glr_darts_002.png
:alt: darts
:srcset: /tutorials/images/sphx_glr_darts_002.png
:class: sphx-glr-single-img
.. GENERATED FROM PYTHON SOURCE LINES 402-412
Retrain the searched model
--------------------------
What we have got in the last step, is only a cell structure.
To get a final usable model with trained weights, we need to construct a real model based on this structure,
and then fully train it.
To construct a fixed model based on the architecture dict exported from the experiment,
we can use :func:`nni.retiarii.fixed_arch`. Seemingly, we are still creating a space.
But under the with-context, we are actually creating a fixed model.
.. GENERATED FROM PYTHON SOURCE LINES 413-419
.. code-block:: default
from nni.retiarii import fixed_arch
with fixed_arch(exported_arch):
final_model = DartsSpace(16, 8, 'cifar')
.. GENERATED FROM PYTHON SOURCE LINES 420-421
We then train the model on full CIFAR-10 training dataset, and evaluate it on the original CIFAR-10 validation dataset.
.. GENERATED FROM PYTHON SOURCE LINES 422-425
.. code-block:: default
train_loader = DataLoader(train_data, batch_size=96, num_workers=6) # Use the original training data
.. GENERATED FROM PYTHON SOURCE LINES 426-427
The validation data loader can be reused.
.. GENERATED FROM PYTHON SOURCE LINES 428-431
.. code-block:: default
valid_loader
.. rst-class:: sphx-glr-script-out
.. code-block:: none
<torch.utils.data.dataloader.DataLoader object at 0x7fba91e9ad60>
.. GENERATED FROM PYTHON SOURCE LINES 432-435
We must create a new evaluator here because a different data split is used.
Also, we should avoid the underlying pytorch-lightning implementation of :class:`~nni.retiarii.evaluator.pytorch.Classification`
evaluator from loading the wrong checkpoint.
.. GENERATED FROM PYTHON SOURCE LINES 436-452
.. code-block:: default
max_epochs = 100
evaluator = Classification(
learning_rate=1e-3,
weight_decay=1e-4,
train_dataloaders=train_loader,
val_dataloaders=valid_loader,
max_epochs=max_epochs,
gpus=1,
export_onnx=False, # Disable ONNX export for this experiment
fast_dev_run=fast_dev_run, # Should be false for fully training
)
evaluator.fit(final_model)
.. rst-class:: sphx-glr-script-out
.. code-block:: none
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:445: LightningDeprecationWarning: Setting `Trainer(gpus=1)` is deprecated in v1.7 and will be removed in v2.0. Please use `Trainer(accelerator='gpu', devices=1)` instead.
rank_zero_deprecation(
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
Running in `fast_dev_run` mode: will run the requested loop using 1 batch(es). Logging and checkpointing is suppressed.
`Trainer(limit_train_batches=1)` was configured so 1 batch per epoch will be used.
`Trainer(limit_val_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_test_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_predict_batches=1)` was configured so 1 batch will be used.
`Trainer(val_check_interval=1.0)` was configured so validation will run at the end of the training epoch..
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
| Name | Type | Params
-----------------------------------------------
0 | criterion | CrossEntropyLoss | 0
1 | metrics | ModuleDict | 0
2 | model | DARTS | 305 K
-----------------------------------------------
305 K Trainable params
0 Non-trainable params
305 K Total params
1.222 Total estimated model params size (MB)
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1891: PossibleUserWarning: The number of training batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
rank_zero_warn(
Training: 0it [00:00, ?it/s] Training: 0%| | 0/2 [00:00<?, ?it/s] Epoch 0: 0%| | 0/2 [00:00<?, ?it/s] Epoch 0: 50%|##### | 1/2 [00:00<00:00, 1.92it/s] Epoch 0: 50%|##### | 1/2 [00:00<00:00, 1.91it/s, loss=2.31, v_num=, train_loss=2.310, train_acc=0.125]
Validation: 0it [00:00, ?it/s]
Validation: 0%| | 0/1 [00:00<?, ?it/s]
Validation DataLoader 0: 0%| | 0/1 [00:00<?, ?it/s]
Validation DataLoader 0: 100%|##########| 1/1 [00:00<00:00, 12.11it/s] Epoch 0: 100%|##########| 2/2 [00:00<00:00, 2.06it/s, loss=2.31, v_num=, train_loss=2.310, train_acc=0.125] Epoch 0: 100%|##########| 2/2 [00:00<00:00, 2.06it/s, loss=2.31, v_num=, train_loss=2.310, train_acc=0.125, val_loss=2.310, val_acc=0.0938]
Epoch 0: 100%|##########| 2/2 [00:00<00:00, 2.05it/s, loss=2.31, v_num=, train_loss=2.310, train_acc=0.125, val_loss=2.310, val_acc=0.0938]`Trainer.fit` stopped: `max_steps=1` reached.
Epoch 0: 100%|##########| 2/2 [00:00<00:00, 2.04it/s, loss=2.31, v_num=, train_loss=2.310, train_acc=0.125, val_loss=2.310, val_acc=0.0938]
.. GENERATED FROM PYTHON SOURCE LINES 453-480
.. note:: When ``fast_dev_run`` is turned off, we achieve a validation accuracy of 89.69% after training for 100 epochs.
Reproduce results in DARTS paper
--------------------------------
After a brief walkthrough of search + retrain process with one-shot strategy,
we then fill the gap between our results (89.69%) and the results in the `DARTS` paper.
This is because we didn't introduce some extra training tricks, including `DropPath <https://arxiv.org/pdf/1605.07648v4.pdf>`__,
Auxiliary loss, gradient clipping and augmentations like `Cutout <https://arxiv.org/pdf/1708.04552v2.pdf>`__.
They also train the deeper (20 cells) and wider (36 channels) networks for longer time (600 epochs).
Evaluator
^^^^^^^^^
To implement these tricks, we first need to rewrite a few parts of evaluator.
Working with one-shot strategies, evaluators need to be implemented in the style of :ref:`PyTorch-Lightning <lightning-evaluator>`,
The full tutorial can be found in :doc:`/nas/evaluator`.
Putting it briefly, the core part of writing a new evaluator is to write a new LightningModule.
`LightingModule <https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html>`__ is a concept in
PyTorch-Lightning, which organizes the model training process into a list of functions, such as,
``training_step``, ``validation_step``, ``configure_optimizers``, etc.
Since we are merely adding a few ingredients to :class:`~nni.retiarii.evaluator.pytorch.Classification`,
we can simply inherit :class:`~nni.retiarii.evaluator.pytorch.ClassificationModule`, which is the underlying LightningModule
behind :class:`~nni.retiarii.evaluator.pytorch.Classification`.
This could look intimidating at first, but most of them are just plug-and-play tricks which you don't need to know details about.
.. GENERATED FROM PYTHON SOURCE LINES 481-536
.. code-block:: default
import torch
from nni.retiarii.evaluator.pytorch import ClassificationModule
class DartsClassificationModule(ClassificationModule):
def __init__(
self,
learning_rate: float = 0.001,
weight_decay: float = 0.,
auxiliary_loss_weight: float = 0.4,
max_epochs: int = 600
):
self.auxiliary_loss_weight = auxiliary_loss_weight
# Training length will be used in LR scheduler
self.max_epochs = max_epochs
super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False)
def configure_optimizers(self):
"""Customized optimizer with momentum, as well as a scheduler."""
optimizer = torch.optim.SGD(
self.parameters(),
momentum=0.9,
lr=self.hparams.learning_rate,
weight_decay=self.hparams.weight_decay
)
return {
'optimizer': optimizer,
'lr_scheduler': torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, self.max_epochs, eta_min=1e-3)
}
def training_step(self, batch, batch_idx):
"""Training step, customized with auxiliary loss."""
x, y = batch
if self.auxiliary_loss_weight:
y_hat, y_aux = self(x)
loss_main = self.criterion(y_hat, y)
loss_aux = self.criterion(y_aux, y)
self.log('train_loss_main', loss_main)
self.log('train_loss_aux', loss_aux)
loss = loss_main + self.auxiliary_loss_weight * loss_aux
else:
y_hat = self(x)
loss = self.criterion(y_hat, y)
self.log('train_loss', loss, prog_bar=True)
for name, metric in self.metrics.items():
self.log('train_' + name, metric(y_hat, y), prog_bar=True)
return loss
def on_train_epoch_start(self):
# Set drop path probability before every epoch. This has no effect if drop path is not enabled in model.
self.model.set_drop_path_prob(self.model.drop_path_prob * self.current_epoch / self.max_epochs)
# Logging learning rate at the beginning of every epoch
self.log('lr', self.trainer.optimizers[0].param_groups[0]['lr'])
.. GENERATED FROM PYTHON SOURCE LINES 537-541
The full evaluator is written as follows,
which simply wraps everything (except model space and search strategy of course), in a single object.
:class:`~nni.retiarii.evaluator.pytorch.Lightning` here is a special type of evaluator.
Don't forget to use the train/val data split specialized for search (1:1) here.
.. GENERATED FROM PYTHON SOURCE LINES 542-558
.. code-block:: default
from nni.retiarii.evaluator.pytorch import Lightning, Trainer
max_epochs = 50
evaluator = Lightning(
DartsClassificationModule(0.025, 3e-4, 0., max_epochs),
Trainer(
gpus=1,
max_epochs=max_epochs,
fast_dev_run=fast_dev_run,
),
train_dataloaders=search_train_loader,
val_dataloaders=search_valid_loader
)
.. rst-class:: sphx-glr-script-out
.. code-block:: none
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:445: LightningDeprecationWarning: Setting `Trainer(gpus=1)` is deprecated in v1.7 and will be removed in v2.0. Please use `Trainer(accelerator='gpu', devices=1)` instead.
rank_zero_deprecation(
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
Running in `fast_dev_run` mode: will run the requested loop using 1 batch(es). Logging and checkpointing is suppressed.
`Trainer(limit_train_batches=1)` was configured so 1 batch per epoch will be used.
`Trainer(limit_val_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_test_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_predict_batches=1)` was configured so 1 batch will be used.
`Trainer(val_check_interval=1.0)` was configured so validation will run at the end of the training epoch..
.. GENERATED FROM PYTHON SOURCE LINES 559-567
Strategy
^^^^^^^^
:class:`~nni.retiarii.strategy.DARTS` strategy is created with gradient clip turned on.
If you are familiar with PyTorch-Lightning, you might aware that gradient clipping can be enabled in Lightning trainer.
However, enabling gradient cip in the trainer above won't work, because the underlying
implementation of :class:`~nni.retiarii.strategy.DARTS` strategy is based on
`manual optimization <https://pytorch-lightning.readthedocs.io/en/stable/common/optimization.html>`__.
.. GENERATED FROM PYTHON SOURCE LINES 568-571
.. code-block:: default
strategy = DartsStrategy(gradient_clip_val=5.)
.. GENERATED FROM PYTHON SOURCE LINES 572-581
Launch experiment
^^^^^^^^^^^^^^^^^
Then we use the newly created evaluator and strategy to launch the experiment again.
.. warning::
``model_space`` has to be re-instantiated because a known limitation,
i.e., one model space can't be reused across multiple experiments.
.. GENERATED FROM PYTHON SOURCE LINES 582-593
.. code-block:: default
model_space = DartsSpace(16, 8, 'cifar')
config = RetiariiExeConfig(execution_engine='oneshot')
experiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)
experiment.run(config)
exported_arch = experiment.export_top_models()[0]
exported_arch
.. rst-class:: sphx-glr-script-out
.. code-block:: none
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
| Name | Type | Params
----------------------------------------------------
0 | model | DartsClassificationModule | 3.0 M
----------------------------------------------------
3.0 M Trainable params
0 Non-trainable params
3.0 M Total params
12.164 Total estimated model params size (MB)
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1891: PossibleUserWarning: The number of training batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
rank_zero_warn(
Training: 0it [00:00, ?it/s] Training: 0%| | 0/1 [00:00<?, ?it/s] Epoch 0: 0%| | 0/1 [00:00<?, ?it/s] Epoch 0: 100%|##########| 1/1 [01:03<00:00, 63.26s/it] Epoch 0: 100%|##########| 1/1 [01:03<00:00, 63.26s/it, v_num=, train_loss=2.420, train_acc=0.0781] Epoch 0: 100%|##########| 1/1 [01:03<00:00, 63.27s/it, v_num=, train_loss=2.420, train_acc=0.0781]`Trainer.fit` stopped: `max_epochs=1` reached.
Epoch 0: 100%|##########| 1/1 [01:03<00:00, 63.27s/it, v_num=, train_loss=2.420, train_acc=0.0781]
{'normal/op_2_0': 'sep_conv_3x3', 'normal/input_2_0': 1, 'normal/op_2_1': 'dil_conv_5x5', 'normal/input_2_1': 0, 'normal/op_3_0': 'dil_conv_5x5', 'normal/input_3_0': 2, 'normal/op_3_1': 'sep_conv_5x5', 'normal/input_3_1': 0, 'normal/op_4_0': 'sep_conv_3x3', 'normal/input_4_0': 2, 'normal/op_4_1': 'sep_conv_3x3', 'normal/input_4_1': 1, 'normal/op_5_0': 'sep_conv_3x3', 'normal/input_5_0': 2, 'normal/op_5_1': 'sep_conv_5x5', 'normal/input_5_1': 3, 'reduce/op_2_0': 'sep_conv_5x5', 'reduce/input_2_0': 1, 'reduce/op_2_1': 'skip_connect', 'reduce/input_2_1': 0, 'reduce/op_3_0': 'sep_conv_3x3', 'reduce/input_3_0': 1, 'reduce/op_3_1': 'dil_conv_3x3', 'reduce/input_3_1': 2, 'reduce/op_4_0': 'sep_conv_3x3', 'reduce/input_4_0': 2, 'reduce/op_4_1': 'avg_pool_3x3', 'reduce/input_4_1': 1, 'reduce/op_5_0': 'dil_conv_3x3', 'reduce/input_5_0': 1, 'reduce/op_5_1': 'sep_conv_3x3', 'reduce/input_5_1': 4}
.. GENERATED FROM PYTHON SOURCE LINES 594-595
We get the following architecture when ``fast_dev_run`` is set to False. It takes around 8 hours on a P100 GPU.
.. GENERATED FROM PYTHON SOURCE LINES 596-632
.. code-block:: default
plot_double_cells({
'normal/op_2_0': 'sep_conv_3x3',
'normal/input_2_0': 0,
'normal/op_2_1': 'sep_conv_3x3',
'normal/input_2_1': 1,
'normal/op_3_0': 'sep_conv_3x3',
'normal/input_3_0': 1,
'normal/op_3_1': 'skip_connect',
'normal/input_3_1': 0,
'normal/op_4_0': 'sep_conv_3x3',
'normal/input_4_0': 0,
'normal/op_4_1': 'max_pool_3x3',
'normal/input_4_1': 1,
'normal/op_5_0': 'sep_conv_3x3',
'normal/input_5_0': 0,
'normal/op_5_1': 'sep_conv_3x3',
'normal/input_5_1': 1,
'reduce/op_2_0': 'max_pool_3x3',
'reduce/input_2_0': 0,
'reduce/op_2_1': 'sep_conv_5x5',
'reduce/input_2_1': 1,
'reduce/op_3_0': 'dil_conv_5x5',
'reduce/input_3_0': 2,
'reduce/op_3_1': 'max_pool_3x3',
'reduce/input_3_1': 0,
'reduce/op_4_0': 'max_pool_3x3',
'reduce/input_4_0': 0,
'reduce/op_4_1': 'sep_conv_3x3',
'reduce/input_4_1': 2,
'reduce/op_5_0': 'max_pool_3x3',
'reduce/input_5_0': 0,
'reduce/op_5_1': 'skip_connect',
'reduce/input_5_1': 2
})
.. image-sg:: /tutorials/images/sphx_glr_darts_003.png
:alt: darts
:srcset: /tutorials/images/sphx_glr_darts_003.png
:class: sphx-glr-single-img
.. GENERATED FROM PYTHON SOURCE LINES 633-640
Retrain
^^^^^^^
When retraining,
we extend the original dataloader to introduce another trick called `Cutout <https://arxiv.org/pdf/1708.04552v2.pdf>`__.
Cutout is a data augmentation technique that randomly masks out rectangular regions in images.
In CIFAR-10, the typical masked size is 16x16 (the image sizes are 32x32 in the dataset).
.. GENERATED FROM PYTHON SOURCE LINES 641-667
.. code-block:: default
def cutout_transform(img, length: int = 16):
h, w = img.size(1), img.size(2)
mask = np.ones((h, w), np.float32)
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - length // 2, 0, h)
y2 = np.clip(y + length // 2, 0, h)
x1 = np.clip(x - length // 2, 0, w)
x2 = np.clip(x + length // 2, 0, w)
mask[y1: y2, x1: x2] = 0.
mask = torch.from_numpy(mask)
mask = mask.expand_as(img)
img *= mask
return img
transform_with_cutout = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
cutout_transform,
])
.. GENERATED FROM PYTHON SOURCE LINES 668-670
The train dataloader needs to be reinstantiated with the new transform.
The validation dataloader is not affected, and thus can be reused.
.. GENERATED FROM PYTHON SOURCE LINES 671-675
.. code-block:: default
train_data_cutout = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform_with_cutout)
train_loader_cutout = DataLoader(train_data_cutout, batch_size=96)
.. rst-class:: sphx-glr-script-out
.. code-block:: none
Files already downloaded and verified
.. GENERATED FROM PYTHON SOURCE LINES 676-681
We then create the final model based on the new exported architecture.
This time, auxiliary loss and drop path probability is enabled.
Following the same procedure as paper, we also increase the number of filters to 36, and number of cells to 20,
so as to reasonably increase the model size and boost the performance.
.. GENERATED FROM PYTHON SOURCE LINES 682-686
.. code-block:: default
with fixed_arch(exported_arch):
final_model = DartsSpace(36, 20, 'cifar', auxiliary_loss=True, drop_path_prob=0.2)
.. GENERATED FROM PYTHON SOURCE LINES 687-688
We create a new evaluator for the retraining process, where the gradient clipping is put into the keyword arguments of trainer.
.. GENERATED FROM PYTHON SOURCE LINES 689-706
.. code-block:: default
max_epochs = 600
evaluator = Lightning(
DartsClassificationModule(0.025, 3e-4, 0.4, max_epochs),
Trainer(
gpus=1,
gradient_clip_val=5.,
max_epochs=max_epochs,
fast_dev_run=fast_dev_run
),
train_dataloaders=train_loader_cutout,
val_dataloaders=valid_loader,
)
evaluator.fit(final_model)
.. rst-class:: sphx-glr-script-out
.. code-block:: none
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:445: LightningDeprecationWarning: Setting `Trainer(gpus=1)` is deprecated in v1.7 and will be removed in v2.0. Please use `Trainer(accelerator='gpu', devices=1)` instead.
rank_zero_deprecation(
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
Running in `fast_dev_run` mode: will run the requested loop using 1 batch(es). Logging and checkpointing is suppressed.
`Trainer(limit_train_batches=1)` was configured so 1 batch per epoch will be used.
`Trainer(limit_val_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_test_batches=1)` was configured so 1 batch will be used.
`Trainer(limit_predict_batches=1)` was configured so 1 batch will be used.
`Trainer(val_check_interval=1.0)` was configured so validation will run at the end of the training epoch..
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
| Name | Type | Params
-----------------------------------------------
0 | criterion | CrossEntropyLoss | 0
1 | metrics | ModuleDict | 0
2 | model | DARTS | 4.8 M
-----------------------------------------------
4.8 M Trainable params
0 Non-trainable params
4.8 M Total params
19.308 Total estimated model params size (MB)
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:219: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 6 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
rank_zero_warn(
/home/yugzhan/miniconda3/envs/nni/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1891: PossibleUserWarning: The number of training batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
rank_zero_warn(
Training: 0it [00:00, ?it/s] Training: 0%| | 0/2 [00:00<?, ?it/s] Epoch 0: 0%| | 0/2 [00:00<?, ?it/s] Epoch 0: 50%|##### | 1/2 [00:00<00:00, 1.15it/s] Epoch 0: 50%|##### | 1/2 [00:00<00:00, 1.15it/s, loss=3.31, v_num=, train_loss=3.310, train_acc=0.115]
Validation: 0it [00:00, ?it/s]
Validation: 0%| | 0/1 [00:00<?, ?it/s]
Validation DataLoader 0: 0%| | 0/1 [00:00<?, ?it/s]
Validation DataLoader 0: 100%|##########| 1/1 [00:00<00:00, 2.84it/s] Epoch 0: 100%|##########| 2/2 [00:01<00:00, 1.27it/s, loss=3.31, v_num=, train_loss=3.310, train_acc=0.115] Epoch 0: 100%|##########| 2/2 [00:01<00:00, 1.27it/s, loss=3.31, v_num=, train_loss=3.310, train_acc=0.115, val_loss=2.300, val_acc=0.113]
Epoch 0: 100%|##########| 2/2 [00:01<00:00, 1.25it/s, loss=3.31, v_num=, train_loss=3.310, train_acc=0.115, val_loss=2.300, val_acc=0.113]`Trainer.fit` stopped: `max_steps=1` reached.
Epoch 0: 100%|##########| 2/2 [00:01<00:00, 1.25it/s, loss=3.31, v_num=, train_loss=3.310, train_acc=0.115, val_loss=2.300, val_acc=0.113]
.. GENERATED FROM PYTHON SOURCE LINES 707-721
When ``fast_dev_run`` is turned off, after retraining, the architecture yields a top-1 accuracy of 97.12%.
If we take the best snapshot throughout the retrain process,
there is a chance that the top-1 accuracy will be 97.28%.
.. image:: ../../img/darts_val_acc.png
In the figure, the orange line is the validation accuracy curve after training for 600 epochs,
while the red line corresponding the previous version in this tutorial before adding all the training tricks and
only trains for 100 epochs.
The results outperforms "DARTS (first order) + cutout" in `DARTS`_ paper, which is only 97.00±0.14%.
It's even comparable with "DARTS (second order) + cutout" in the paper (97.24±0.09%),
though we didn't implement the second order version.
The implementation of second order DARTS is in our future plan, and we also welcome your contribution.
.. rst-class:: sphx-glr-timing
**Total running time of the script:** ( 1 minutes 38.004 seconds)
.. _sphx_glr_download_tutorials_darts.py:
.. only:: html
.. container:: sphx-glr-footer sphx-glr-footer-example
.. container:: sphx-glr-download sphx-glr-download-python
:download:`Download Python source code: darts.py <darts.py>`
.. container:: sphx-glr-download sphx-glr-download-jupyter
:download:`Download Jupyter notebook: darts.ipynb <darts.ipynb>`
.. only:: html
.. rst-class:: sphx-glr-signature
`Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_
docs/source/tutorials/index.rst
View file @ bddaf36d
......@@ -146,6 +146,23 @@ Tutorials
</div>
.. raw:: html
<div class="sphx-glr-thumbcontainer" tooltip="In this tutorial, we demonstrate how to search in the famous model space proposed in `DARTS`_.">
.. only:: html
.. image:: /tutorials/images/thumb/sphx_glr_darts_thumb.png
:alt: Searching in DARTS search space
:ref:`sphx_glr_tutorials_darts.py`
.. raw:: html
<div class="sphx-glr-thumbnail-title">Searching in DARTS search space</div>
</div>
.. raw:: html
<div class="sphx-glr-thumbcontainer" tooltip="Workable Pruning Process ------------------------">
......@@ -179,6 +196,7 @@ Tutorials
/tutorials/nasbench_as_dataset
/tutorials/pruning_customize
/tutorials/hello_nas
/tutorials/darts
/tutorials/pruning_bert_glue
......
docs/source/tutorials/sg_execution_times.rst
View file @ bddaf36d
......@@ -5,8 +5,10 @@
Computation times
=================
**00:27.206** total execution time for **tutorials** files:
**01:38.004** total execution time for **tutorials** files:
+-----------------------------------------------------------------------------------------------------+-----------+--------+
| :ref:`sphx_glr_tutorials_darts.py` (``darts.py``) | 01:38.004 | 0.0 MB |
+-----------------------------------------------------------------------------------------------------+-----------+--------+
| :ref:`sphx_glr_tutorials_pruning_bert_glue.py` (``pruning_bert_glue.py``) | 00:27.206 | 0.0 MB |
+-----------------------------------------------------------------------------------------------------+-----------+--------+
......
examples/tutorials/.gitignore
View file @ bddaf36d
data/
log/
*.onnx
lightning_logs
models/
pruning_log/
examples/tutorials/darts.py 0 → 100644
View file @ bddaf36d
"""
Searching in DARTS search space
===============================
In this tutorial, we demonstrate how to search in the famous model space proposed in `DARTS`_.
Through this process, you will learn:
* How to use the built-in model spaces from NNI's model space hub.
* How to use one-shot exploration strategies to explore a model space.
* How to customize evaluators to achieve the best performance.
In the end, we get a strong-performing model on CIFAR-10 dataset, which achieves up to 97.28% accuracy.
.. attention::
Running this tutorial requires a GPU.
If you don't have one, you can set ``gpus`` in :class:`~nni.retiarii.evaluator.pytorch.Classification` to be 0,
but do note that it will be much slower.
.. _DARTS: https://arxiv.org/abs/1806.09055
Use a pre-searched model
------------------------
Similar to `the beginner tutorial of PyTorch <https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html>`__,
we begin with CIFAR-10 dataset, which is a image classification dataset of 10 categories.
The images in CIFAR-10 are of size 3x32x32, i.e., RGB-colored images of 32x32 pixels in size.
We first load the CIFAR-10 dataset with torchvision.
"""
import nni
import torch
from torchvision import transforms
from torchvision.datasets import CIFAR10
from nni.retiarii.evaluator.pytorch import DataLoader
CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]
transform_valid = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
valid_data = nni.trace(CIFAR10)(root='./data', train=False, download=True, transform=transform_valid)
valid_loader = DataLoader(valid_data, batch_size=256, num_workers=6)
# %%
#
# .. note::
#
# If you are to use multi-trial strategies, wrapping CIFAR10 with :func:`nni.trace` and
# use DataLoader from ``nni.retiarii.evaluator.pytorch`` (instead of ``torch.utils.data``) are mandatory.
# Otherwise, it's optional.
#
# When working with famous datasets like CIFAR-10 or ImageNet,
# it's tempting to use or finetune from a pretrained model, like ResNet.
# There's nothing wrong with doing so, and sometimes it might be beneficial.
# Thanks to the development of NAS, we now have quite a large number of *pre-searched models*,
# which are produced by most popular NAS literatures.
# You can easily load these models, validate their performances, and finetune them if you need.
#
# We present :doc:`model space hub </nas/space_hub>`, where you can find many built-in model spaces,
# along with many pre-searched models.
# We choose one from `DARTS`_ search space, which is natively trained on our target dataset, CIFAR-10,
# so as to save the tedious steps of finetuning.
#
# .. tip::
#
# Finetuning a pre-searched model on other datasets is no different from finetuning *any model*.
# We recommend reading
# `this tutorial of object detection finetuning <https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html>`__
# if you want to know how finetuning is generally done in PyTorch.
from nni.retiarii.hub.pytorch import DARTS as DartsSpace
darts_v2_model = DartsSpace.load_searched_model('darts-v2', pretrained=True, download=True)
def evaluate_model(model, cuda=False):
device = torch.device('cuda' if cuda else 'cpu')
model.to(device)
model.eval()
with torch.no_grad():
correct = total = 0
for inputs, targets in valid_loader:
inputs, targets = inputs.to(device), targets.to(device)
logits = model(inputs)
_, predict = torch.max(logits, 1)
correct += (predict == targets).sum().cpu().item()
total += targets.size(0)
print('Accuracy:', correct / total)
return correct / total
evaluate_model(darts_v2_model, True) # Set this to false if there's no GPU.
# %%
#
# The journey could end here. Or you are interested,
# we can go a step further to search a model within :class:`~nni.retiarii.hub.pytorch.DARTS` space on our own.
#
# Use the model space
# -------------------
#
# The model space provided in `DARTS`_ originated from `NASNet <https://arxiv.org/abs/1707.07012>`__,
# where the full model is constructed by repeatedly stacking a single computational unit (called a **cell**).
# There are two types of cells within a network. The first type is called *normal cell*, and the second type is called *reduction cell*.
# The key difference between normal and reduction cell is that the reduction cell will downsample the input feature map,
# and decrease its resolution. Normal and reduction cells are stacked alternately, as shown in the following figure.
#
# .. image:: ../../img/nasnet_cell_stack.png
#
# A cell takes outputs from two previous cells as inputs and contains a collection of *nodes*.
# Each node takes two previous nodes within the same cell (or the two cell inputs),
# and applies an *operator* (e.g., convolution, or max-pooling) to each input,
# and sums the outputs of operators as the output of the node.
# The output of cell is the concatenation of all the nodes that are never used as inputs of another node.
# We recommend reading `NDS <https://arxiv.org/pdf/1905.13214.pdf>`__ or `ENAS <https://arxiv.org/abs/1802.03268>`__ for details.
#
# We illustrate an example of cells in the following figure.
#
# .. image:: ../../img/nasnet_cell.png
#
# The search space proposed in `DARTS`_ paper introduced two modifications to the original space
# in `NASNet <https://arxiv.org/abs/1707.07012>`__.
#
# Firstly, the operator candidates have been narrowed down to seven:
#
# - Max pooling 3x3
# - Average pooling 3x3
# - Skip connect (Identity)
# - Separable convolution 3x3
# - Separable convolution 5x5
# - Dilated convolution 3x3
# - Dilated convolution 5x5
#
# Secondly, the output of cell is the concatenate of **all the nodes within the cell**.
#
# As the search space is based on cell, once the normal and reduction cell has been fixed, we can stack them for indefinite times.
# To save the search cost, the common practice is to reduce the number of filters (i.e., channels) and number of stacked cells
# during the search phase, and increase them back when training the final searched architecture.
#
# .. note::
#
# `DARTS`_ is one of those papers that innovate both in search space and search strategy.
# In this tutorial, we will search on **model space** provided by DARTS with **search strategy** proposed by DARTS.
# We refer to them as *DARTS model space* (``DartsSpace``) and *DARTS strategy* (``DartsStrategy``), respectively.
# We did NOT imply that the :class:`~nni.retiarii.hub.pytorch.DARTS` space and
# :class:`~nni.retiarii.strategy.DARTS` strategy has to used together.
# You can always explore the DARTS space with another search strategy, or use your own strategy to search a different model space.
#
# In the following example, we initialize a :class:`~nni.retiarii.hub.pytorch.DARTS`
# model space, with 16 initial filters and 8 stacked cells.
# The network is specialized for CIFAR-10 dataset with 32x32 input resolution.
#
# The :class:`~nni.retiarii.hub.pytorch.DARTS` model space here is provided by :doc:`model space hub </nas/space_hub>`,
# where we have supported multiple popular model spaces for plug-and-play.
#
# .. tip::
#
# The model space here can be replaced with any space provided in the hub,
# or even customized spaces built from scratch.
model_space = DartsSpace(16, 8, 'cifar')
# %%
#
# Search on the model space
# -------------------------
#
# .. warning::
#
# Please set ``fast_dev_run`` to False to reproduce the our claimed results.
# Otherwise, only a few mini-batches will be run.
fast_dev_run = True
# %%
#
# Evaluator
# ^^^^^^^^^
#
# To begin exploring the model space, one firstly need to have an evaluator to provide the criterion of a "good model".
# As we are searching on CIFAR-10 dataset, one can easily use the :class:`~nni.retiarii.evaluator.pytorch.Classification`
# as a starting point.
#
# Note that for a typical setup of NAS, the model search should be on validation set, and the evaluation of the final searched model
# should be on test set. However, as CIFAR-10 dataset doesn't have a test dataset (only 50k train + 10k valid),
# we have to split the original training set into a training set and a validation set.
# The recommended train/val split by `DARTS`_ strategy is 1:1.
import numpy as np
from nni.retiarii.evaluator.pytorch import Classification
from torch.utils.data import SubsetRandomSampler
transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
train_data = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform)
num_samples = len(train_data)
indices = np.random.permutation(num_samples)
split = num_samples // 2
search_train_loader = DataLoader(
train_data, batch_size=64, num_workers=6,
sampler=SubsetRandomSampler(indices[:split]),
)
search_valid_loader = DataLoader(
train_data, batch_size=64, num_workers=6,
sampler=SubsetRandomSampler(indices[split:]),
)
evaluator = Classification(
learning_rate=1e-3,
weight_decay=1e-4,
train_dataloaders=search_train_loader,
val_dataloaders=search_valid_loader,
max_epochs=10,
gpus=1,
fast_dev_run=fast_dev_run,
)
# %%
#
# Strategy
# ^^^^^^^^
#
# We will use `DARTS`_ (Differentiable ARchiTecture Search) as the search strategy to explore the model space.
# :class:`~nni.retiarii.strategy.DARTS` strategy belongs to the category of :ref:`one-shot strategy <one-shot-nas>`.
# The fundamental differences between One-shot strategies and :ref:`multi-trial strategies <multi-trial-nas>` is that,
# one-shot strategy combines search with model training into a single run.
# Compared to multi-trial strategies, one-shot NAS doesn't need to iteratively spawn new trials (i.e., models),
# and thus saves the excessive cost of model training.
# It's worth mentioning that one-shot NAS also suffers from multiple drawbacks despite its computational efficiency.
# We recommend
# `Weight-Sharing Neural Architecture Search: A Battle to Shrink the Optimization Gap <https://arxiv.org/abs/2008.01475>`__
# and
# `How Does Supernet Help in Neural Architecture Search? <https://arxiv.org/abs/2010.08219>`__ for interested readers.
#
# :class:`~nni.retiarii.strategy.DARTS` strategy is provided as one of NNI's :doc:`built-in search strategies </nas/exploration_strategy>`.
# Using it can be as simple as one line of code.
from nni.retiarii.strategy import DARTS as DartsStrategy
strategy = DartsStrategy()
# %%
#
# .. tip:: The ``DartsStrategy`` here can be replaced by any search strategies, even multi-trial strategies.
#
# If you want to know how DARTS strategy works, here is a brief version.
# Under the hood, DARTS converts the cell into a densely connected graph, and put operators on edges (see the following figure).
# Since the operators are not decided yet, every edge is a weighted mixture of multiple operators (multiple color in the figure).
# DARTS then learns to assign the optimal "color" for each edge during the network training.
# It finally selects one "color" for each edge, and drops redundant edges.
# The weights on the edges are called *architecture weights*.
#
# .. image:: ../../img/darts_illustration.png
#
# It's NOT reflected in the figure that, for DARTS model space, exactly two inputs are kept for every node.
#
# Launch experiment
# ^^^^^^^^^^^^^^^^^
#
# We then come to the step of launching the experiment.
# This step is similar to what we have done in the :doc:`beginner tutorial <hello_nas>`,
# except that the ``execution_engine`` argument should be set to ``oneshot``.
from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
config = RetiariiExeConfig(execution_engine='oneshot')
experiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)
experiment.run(config)
# %%
#
# .. tip::
#
# The search process can be visualized with tensorboard. For example::
#
# tensorboard --logdir=./lightning_logs
#
# Then, open the browser and go to http://localhost:6006/ to monitor the search process.
#
# .. image:: ../../img/darts_search_process.png
#
# We can then retrieve the best model found by the strategy with ``export_top_models``.
# Here, the retrieved model is a dict (called *architecture dict*) describing the selected normal cell and reduction cell.
exported_arch = experiment.export_top_models()[0]
exported_arch
# %%
#
# The cell can be visualized with the following code snippet
# (copied and modified from `DARTS visualization <https://github.com/quark0/darts/blob/master/cnn/visualize.py>`__).
import io
import graphviz
import matplotlib.pyplot as plt
from PIL import Image
def plot_single_cell(arch_dict, cell_name):
g = graphviz.Digraph(
node_attr=dict(style='filled', shape='rect', align='center'),
format='png'
)
g.body.extend(['rankdir=LR'])
g.node('c_{k-2}', fillcolor='darkseagreen2')
g.node('c_{k-1}', fillcolor='darkseagreen2')
assert len(arch_dict) % 2 == 0
for i in range(2, 6):
g.node(str(i), fillcolor='lightblue')
for i in range(2, 6):
for j in range(2):
op = arch_dict[f'{cell_name}/op_{i}_{j}']
from_ = arch_dict[f'{cell_name}/input_{i}_{j}']
if from_ == 0:
u = 'c_{k-2}'
elif from_ == 1:
u = 'c_{k-1}'
else:
u = str(from_)
v = str(i)
g.edge(u, v, label=op, fillcolor='gray')
g.node('c_{k}', fillcolor='palegoldenrod')
for i in range(2, 6):
g.edge(str(i), 'c_{k}', fillcolor='gray')
g.attr(label=f'{cell_name.capitalize()} cell')
image = Image.open(io.BytesIO(g.pipe()))
return image
def plot_double_cells(arch_dict):
image1 = plot_single_cell(arch_dict, 'normal')
image2 = plot_single_cell(arch_dict, 'reduce')
height_ratio = max(image1.size[1] / image1.size[0], image2.size[1] / image2.size[0])
_, axs = plt.subplots(1, 2, figsize=(20, 10 * height_ratio))
axs[0].imshow(image1)
axs[1].imshow(image2)
axs[0].axis('off')
axs[1].axis('off')
plt.show()
plot_double_cells(exported_arch)
# %%
#
# .. warning:: The cell above is obtained via ``fast_dev_run`` (i.e., running only 1 mini-batch).
#
# When ``fast_dev_run`` is turned off, we get a model with the following architecture,
# where you might notice an interesting fact that around half the operations have selected ``sep_conv_3x3``.
plot_double_cells({
'normal/op_2_0': 'sep_conv_3x3',
'normal/input_2_0': 1,
'normal/op_2_1': 'sep_conv_3x3',
'normal/input_2_1': 0,
'normal/op_3_0': 'sep_conv_3x3',
'normal/input_3_0': 1,
'normal/op_3_1': 'sep_conv_3x3',
'normal/input_3_1': 2,
'normal/op_4_0': 'sep_conv_3x3',
'normal/input_4_0': 1,
'normal/op_4_1': 'sep_conv_3x3',
'normal/input_4_1': 0,
'normal/op_5_0': 'sep_conv_3x3',
'normal/input_5_0': 1,
'normal/op_5_1': 'max_pool_3x3',
'normal/input_5_1': 0,
'reduce/op_2_0': 'sep_conv_3x3',
'reduce/input_2_0': 0,
'reduce/op_2_1': 'sep_conv_3x3',
'reduce/input_2_1': 1,
'reduce/op_3_0': 'dil_conv_5x5',
'reduce/input_3_0': 2,
'reduce/op_3_1': 'sep_conv_3x3',
'reduce/input_3_1': 0,
'reduce/op_4_0': 'dil_conv_5x5',
'reduce/input_4_0': 2,
'reduce/op_4_1': 'sep_conv_5x5',
'reduce/input_4_1': 1,
'reduce/op_5_0': 'sep_conv_5x5',
'reduce/input_5_0': 4,
'reduce/op_5_1': 'dil_conv_5x5',
'reduce/input_5_1': 2
})
# %%
#
# Retrain the searched model
# --------------------------
#
# What we have got in the last step, is only a cell structure.
# To get a final usable model with trained weights, we need to construct a real model based on this structure,
# and then fully train it.
#
# To construct a fixed model based on the architecture dict exported from the experiment,
# we can use :func:`nni.retiarii.fixed_arch`. Seemingly, we are still creating a space.
# But under the with-context, we are actually creating a fixed model.
from nni.retiarii import fixed_arch
with fixed_arch(exported_arch):
final_model = DartsSpace(16, 8, 'cifar')
# %%
#
# We then train the model on full CIFAR-10 training dataset, and evaluate it on the original CIFAR-10 validation dataset.
train_loader = DataLoader(train_data, batch_size=96, num_workers=6) # Use the original training data
# %%
#
# The validation data loader can be reused.
valid_loader
# %%
#
# We must create a new evaluator here because a different data split is used.
# Also, we should avoid the underlying pytorch-lightning implementation of :class:`~nni.retiarii.evaluator.pytorch.Classification`
# evaluator from loading the wrong checkpoint.
max_epochs = 100
evaluator = Classification(
learning_rate=1e-3,
weight_decay=1e-4,
train_dataloaders=train_loader,
val_dataloaders=valid_loader,
max_epochs=max_epochs,
gpus=1,
export_onnx=False, # Disable ONNX export for this experiment
fast_dev_run=fast_dev_run, # Should be false for fully training
)
evaluator.fit(final_model)
# %%
#
# .. note:: When ``fast_dev_run`` is turned off, we achieve a validation accuracy of 89.69% after training for 100 epochs.
#
# Reproduce results in DARTS paper
# --------------------------------
#
# After a brief walkthrough of search + retrain process with one-shot strategy,
# we then fill the gap between our results (89.69%) and the results in the `DARTS` paper.
# This is because we didn't introduce some extra training tricks, including `DropPath <https://arxiv.org/pdf/1605.07648v4.pdf>`__,
# Auxiliary loss, gradient clipping and augmentations like `Cutout <https://arxiv.org/pdf/1708.04552v2.pdf>`__.
# They also train the deeper (20 cells) and wider (36 channels) networks for longer time (600 epochs).
#
#
# Evaluator
# ^^^^^^^^^
#
# To implement these tricks, we first need to rewrite a few parts of evaluator.
#
# Working with one-shot strategies, evaluators need to be implemented in the style of :ref:`PyTorch-Lightning <lightning-evaluator>`,
# The full tutorial can be found in :doc:`/nas/evaluator`.
# Putting it briefly, the core part of writing a new evaluator is to write a new LightningModule.
# `LightingModule <https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html>`__ is a concept in
# PyTorch-Lightning, which organizes the model training process into a list of functions, such as,
# ``training_step``, ``validation_step``, ``configure_optimizers``, etc.
# Since we are merely adding a few ingredients to :class:`~nni.retiarii.evaluator.pytorch.Classification`,
# we can simply inherit :class:`~nni.retiarii.evaluator.pytorch.ClassificationModule`, which is the underlying LightningModule
# behind :class:`~nni.retiarii.evaluator.pytorch.Classification`.
# This could look intimidating at first, but most of them are just plug-and-play tricks which you don't need to know details about.
import torch
from nni.retiarii.evaluator.pytorch import ClassificationModule
class DartsClassificationModule(ClassificationModule):
def __init__(
self,
learning_rate: float = 0.001,
weight_decay: float = 0.,
auxiliary_loss_weight: float = 0.4,
max_epochs: int = 600
):
self.auxiliary_loss_weight = auxiliary_loss_weight
# Training length will be used in LR scheduler
self.max_epochs = max_epochs
super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False)
def configure_optimizers(self):
"""Customized optimizer with momentum, as well as a scheduler."""
optimizer = torch.optim.SGD(
self.parameters(),
momentum=0.9,
lr=self.hparams.learning_rate,
weight_decay=self.hparams.weight_decay
)
return {
'optimizer': optimizer,
'lr_scheduler': torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, self.max_epochs, eta_min=1e-3)
}
def training_step(self, batch, batch_idx):
"""Training step, customized with auxiliary loss."""
x, y = batch
if self.auxiliary_loss_weight:
y_hat, y_aux = self(x)
loss_main = self.criterion(y_hat, y)
loss_aux = self.criterion(y_aux, y)
self.log('train_loss_main', loss_main)
self.log('train_loss_aux', loss_aux)
loss = loss_main + self.auxiliary_loss_weight * loss_aux
else:
y_hat = self(x)
loss = self.criterion(y_hat, y)
self.log('train_loss', loss, prog_bar=True)
for name, metric in self.metrics.items():
self.log('train_' + name, metric(y_hat, y), prog_bar=True)
return loss
def on_train_epoch_start(self):
# Set drop path probability before every epoch. This has no effect if drop path is not enabled in model.
self.model.set_drop_path_prob(self.model.drop_path_prob * self.current_epoch / self.max_epochs)
# Logging learning rate at the beginning of every epoch
self.log('lr', self.trainer.optimizers[0].param_groups[0]['lr'])
# %%
#
# The full evaluator is written as follows,
# which simply wraps everything (except model space and search strategy of course), in a single object.
# :class:`~nni.retiarii.evaluator.pytorch.Lightning` here is a special type of evaluator.
# Don't forget to use the train/val data split specialized for search (1:1) here.
from nni.retiarii.evaluator.pytorch import Lightning, Trainer
max_epochs = 50
evaluator = Lightning(
DartsClassificationModule(0.025, 3e-4, 0., max_epochs),
Trainer(
gpus=1,
max_epochs=max_epochs,
fast_dev_run=fast_dev_run,
),
train_dataloaders=search_train_loader,
val_dataloaders=search_valid_loader
)
# %%
#
# Strategy
# ^^^^^^^^
#
# :class:`~nni.retiarii.strategy.DARTS` strategy is created with gradient clip turned on.
# If you are familiar with PyTorch-Lightning, you might aware that gradient clipping can be enabled in Lightning trainer.
# However, enabling gradient cip in the trainer above won't work, because the underlying
# implementation of :class:`~nni.retiarii.strategy.DARTS` strategy is based on
# `manual optimization <https://pytorch-lightning.readthedocs.io/en/stable/common/optimization.html>`__.
strategy = DartsStrategy(gradient_clip_val=5.)
# %%
#
# Launch experiment
# ^^^^^^^^^^^^^^^^^
#
# Then we use the newly created evaluator and strategy to launch the experiment again.
#
# .. warning::
#
# ``model_space`` has to be re-instantiated because a known limitation,
# i.e., one model space can't be reused across multiple experiments.
model_space = DartsSpace(16, 8, 'cifar')
config = RetiariiExeConfig(execution_engine='oneshot')
experiment = RetiariiExperiment(model_space, evaluator=evaluator, strategy=strategy)
experiment.run(config)
exported_arch = experiment.export_top_models()[0]
exported_arch
# %%
#
# We get the following architecture when ``fast_dev_run`` is set to False. It takes around 8 hours on a P100 GPU.
plot_double_cells({
'normal/op_2_0': 'sep_conv_3x3',
'normal/input_2_0': 0,
'normal/op_2_1': 'sep_conv_3x3',
'normal/input_2_1': 1,
'normal/op_3_0': 'sep_conv_3x3',
'normal/input_3_0': 1,
'normal/op_3_1': 'skip_connect',
'normal/input_3_1': 0,
'normal/op_4_0': 'sep_conv_3x3',
'normal/input_4_0': 0,
'normal/op_4_1': 'max_pool_3x3',
'normal/input_4_1': 1,
'normal/op_5_0': 'sep_conv_3x3',
'normal/input_5_0': 0,
'normal/op_5_1': 'sep_conv_3x3',
'normal/input_5_1': 1,
'reduce/op_2_0': 'max_pool_3x3',
'reduce/input_2_0': 0,
'reduce/op_2_1': 'sep_conv_5x5',
'reduce/input_2_1': 1,
'reduce/op_3_0': 'dil_conv_5x5',
'reduce/input_3_0': 2,
'reduce/op_3_1': 'max_pool_3x3',
'reduce/input_3_1': 0,
'reduce/op_4_0': 'max_pool_3x3',
'reduce/input_4_0': 0,
'reduce/op_4_1': 'sep_conv_3x3',
'reduce/input_4_1': 2,
'reduce/op_5_0': 'max_pool_3x3',
'reduce/input_5_0': 0,
'reduce/op_5_1': 'skip_connect',
'reduce/input_5_1': 2
})
# %%
#
# Retrain
# ^^^^^^^
#
# When retraining,
# we extend the original dataloader to introduce another trick called `Cutout <https://arxiv.org/pdf/1708.04552v2.pdf>`__.
# Cutout is a data augmentation technique that randomly masks out rectangular regions in images.
# In CIFAR-10, the typical masked size is 16x16 (the image sizes are 32x32 in the dataset).
def cutout_transform(img, length: int = 16):
h, w = img.size(1), img.size(2)
mask = np.ones((h, w), np.float32)
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - length // 2, 0, h)
y2 = np.clip(y + length // 2, 0, h)
x1 = np.clip(x - length // 2, 0, w)
x2 = np.clip(x + length // 2, 0, w)
mask[y1: y2, x1: x2] = 0.
mask = torch.from_numpy(mask)
mask = mask.expand_as(img)
img *= mask
return img
transform_with_cutout = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
cutout_transform,
])
# %%
#
# The train dataloader needs to be reinstantiated with the new transform.
# The validation dataloader is not affected, and thus can be reused.
train_data_cutout = nni.trace(CIFAR10)(root='./data', train=True, download=True, transform=transform_with_cutout)
train_loader_cutout = DataLoader(train_data_cutout, batch_size=96)
# %%
#
# We then create the final model based on the new exported architecture.
# This time, auxiliary loss and drop path probability is enabled.
#
# Following the same procedure as paper, we also increase the number of filters to 36, and number of cells to 20,
# so as to reasonably increase the model size and boost the performance.
with fixed_arch(exported_arch):
final_model = DartsSpace(36, 20, 'cifar', auxiliary_loss=True, drop_path_prob=0.2)
# %%
#
# We create a new evaluator for the retraining process, where the gradient clipping is put into the keyword arguments of trainer.
max_epochs = 600
evaluator = Lightning(
DartsClassificationModule(0.025, 3e-4, 0.4, max_epochs),
Trainer(
gpus=1,
gradient_clip_val=5.,
max_epochs=max_epochs,
fast_dev_run=fast_dev_run
),
train_dataloaders=train_loader_cutout,
val_dataloaders=valid_loader,
)
evaluator.fit(final_model)
# %%
#
# When ``fast_dev_run`` is turned off, after retraining, the architecture yields a top-1 accuracy of 97.12%.
# If we take the best snapshot throughout the retrain process,
# there is a chance that the top-1 accuracy will be 97.28%.
#
# .. image:: ../../img/darts_val_acc.png
#
# In the figure, the orange line is the validation accuracy curve after training for 600 epochs,
# while the red line corresponding the previous version in this tutorial before adding all the training tricks and
# only trains for 100 epochs.
#
# The results outperforms "DARTS (first order) + cutout" in `DARTS`_ paper, which is only 97.00±0.14%.
# It's even comparable with "DARTS (second order) + cutout" in the paper (97.24±0.09%),
# though we didn't implement the second order version.
# The implementation of second order DARTS is in our future plan, and we also welcome your contribution.
nni/nas/evaluator/pytorch/lightning.py
View file @ bddaf36d
......@@ -260,12 +260,12 @@ class SupervisedLearningModule(LightningModule):
return self.optimizer(self.parameters(), lr=self.hparams.learning_rate, weight_decay=self.hparams.weight_decay) # type: ignore
def on_validation_epoch_end(self):
if not self.trainer.sanity_checking and self.running_mode == 'multi':
if not self.trainer.sanity_checking and self.running_mode == 'multi' and nni.get_current_parameter() is not None:
# Don't report metric when sanity checking
nni.report_intermediate_result(self._get_validation_metrics())
def on_fit_end(self):
if self.running_mode == 'multi':
if self.running_mode == 'multi' and nni.get_current_parameter() is not None:
nni.report_final_result(self._get_validation_metrics())
def _get_validation_metrics(self):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment