`SPTAG <https://github.com/microsoft/SPTAG>`__ (Space Partition Tree And Graph) is a library for large scale vector approximate nearest neighbor search scenario released by `Microsoft Research (MSR) <https://www.msra.cn/>`__ and `Microsoft Bing <https://www.bing.com/>`__.
This library assumes that the samples are represented as vectors and that the vectors can be compared by L2 distances or cosine distances. Vectors returned for a query vector are the vectors that have smallest L2 distance or cosine distances with the query vector.
SPTAG provides two methods: kd-tree and relative neighborhood graph (SPTAG-KDT) and balanced k-means tree and relative neighborhood graph (SPTAG-BKT). SPTAG-KDT is advantageous in index building cost, and SPTAG-BKT is advantageous in search accuracy in very high-dimensional data.
In SPTAG, there are tens of parameters that can be tuned for specified scenarios or datasets. NNI is a great tool for automatically tuning those parameters. The authors of SPTAG tried NNI for the auto tuning and found good-performing parameters easily, thus, they shared the practice of tuning SPTAG on NNI in their document `here <https://github.com/microsoft/SPTAG/blob/master/docs/Parameters.md>`__. Please refer to it for detailed tutorial.
In the **trial** part, if you want to use GPU to perform the architecture search, change ``trialGpuNum`` from ``0`` to ``1``. You need to increase the ``maxTrialNumber`` and ``maxExperimentDuration``\ , according to how long you want to wait for the search result.
As we can see, this function is actually a compiler, that converts the internal model DAG configuration (which will be introduced in the ``Model configuration format`` section) ``graph``\ , to a Tensorflow computation graph.
.. code-block:: python
topology = graph.is_topology()
performs topological sorting on the internal graph representation, and the code inside the loop:
.. code-block:: python
for _, topo_i in enumerate(topology):
...
performs actually conversion that maps each layer to a part in Tensorflow computation graph.
3.3 The tuner
^^^^^^^^^^^^^
The tuner is much more simple than the trial. They actually share the same ``graph.py``. Besides, the tuner has a ``customer_tuner.py``\ , the most important class in which is ``CustomerTuner``\ :
.. code-block:: python
class CustomerTuner(Tuner):
# ......
def generate_parameters(self, parameter_id):
"""Returns a set of trial graph config, as a serializable object.
parameter_id : int
"""
if len(self.population) <= 0:
logger.debug("the len of poplution lower than zero.")
raise Exception('The population is empty')
pos = -1
for i in range(len(self.population)):
if self.population[i].result == None:
pos = i
break
if pos != -1:
indiv = copy.deepcopy(self.population[pos])
self.population.pop(pos)
temp = json.loads(graph_dumps(indiv.config))
else:
random.shuffle(self.population)
if self.population[0].result > self.population[1].result:
self.population[0] = self.population[1]
indiv = copy.deepcopy(self.population[0])
self.population.pop(1)
indiv.mutation()
graph = indiv.config
temp = json.loads(graph_dumps(graph))
# ......
As we can see, the overloaded method ``generate_parameters`` implements a pretty naive mutation algorithm. The code lines:
.. code-block:: python
if self.population[0].result > self.population[1].result:
self.population[0] = self.population[1]
indiv = copy.deepcopy(self.population[0])
controls the mutation process. It will always take two random individuals in the population, only keeping and mutating the one with better result.
3.4 Model configuration format
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Here is an example of the model configuration, which is passed from the tuner to the trial in the architecture search procedure.
.. code-block:: json
{
"max_layer_num": 50,
"layers": [
{
"input_size": 0,
"type": 3,
"output_size": 1,
"input": [],
"size": "x",
"output": [4, 5],
"is_delete": false
},
{
"input_size": 0,
"type": 3,
"output_size": 1,
"input": [],
"size": "y",
"output": [4, 5],
"is_delete": false
},
{
"input_size": 1,
"type": 4,
"output_size": 0,
"input": [6],
"size": "x",
"output": [],
"is_delete": false
},
{
"input_size": 1,
"type": 4,
"output_size": 0,
"input": [5],
"size": "y",
"output": [],
"is_delete": false
},
{"Comment": "More layers will be here for actual graphs."}
]
}
Every model configuration will have a "layers" section, which is a JSON list of layer definitions. The definition of each layer is also a JSON object, where:
* ``type`` is the type of the layer. 0, 1, 2, 3, 4 corresponds to attention, self-attention, RNN, input and output layer respectively.
* ``size`` is the length of the output. "x", "y" correspond to document length / question length, respectively.
* ``input_size`` is the number of inputs the layer has.
* ``input`` is the indices of layers taken as input of this layer.
* ``output`` is the indices of layers use this layer's output as their input.
* ``is_delete`` means whether the layer is still available.
"\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 DARTS 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"
"<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\nNNI presents many built-in model spaces, along with many *pre-searched models* in :doc:`model space hub </nas/space_hub>`,\nwhich are produced by most popular NAS literatures.\nA pre-trained model is a saved network that was previously trained on a large dataset like CIFAR-10 or ImageNet.\nYou can easily load these models as a starting point, validate their performances, and finetune them if you need.\n\nIn this tutorial, we 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, cuda=True) # Set this to false if there's no GPU."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The journey of using a pre-searched model 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 DARTS 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.\nUsers could read [NDS](https://arxiv.org/pdf/1905.13214.pdf)_ or [ENAS](https://arxiv.org/abs/1802.03268)_ for more 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(\n width=16, # the initial filters (channel number) for the model\n num_cells=8, # the number of stacked cells in total\n dataset='cifar' # to give a hint about input resolution, here is 32x32\n)"
]
},
{
"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"
"### 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.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>It's worth mentioning that one-shot NAS also suffers from multiple drawbacks despite its computational efficiency.\n We recommend\n [Weight-Sharing Neural Architecture Search: A Battle to Shrink the Optimization Gap](https://arxiv.org/abs/2008.01475)_\n and\n [How Does Supernet Help in Neural Architecture Search?](https://arxiv.org/abs/2010.08219)_ for interested readers.</p></div>\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\n.. tip:: It'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"
".. 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"
"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"
"## 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`. Under the with-context, we will creating a fixed model based on ``exported_arch``,\ninstead of creating a space.\n\n"
"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 filters) networks for longer time (600 epochs).\nHere we reproduce these tricks to get comparable results with DARTS paper.\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"
"### 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 clip 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 instance can't be reused across multiple experiments.</p></div>\n\n"
"### 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"
"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"
"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"
/data/data0/jiahang/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:447: 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.
.. GENERATED FROM PYTHON SOURCE LINES 230-250
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'tneedtoiterativelyspawnnewtrials(i.e.,models),
andthussavestheexcessivecostofmodeltraining.
..note::
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-insearchstrategies</nas/exploration_strategy>`.
/data/data0/jiahang/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1892: 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.
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 425-428
.. code-block:: default
train_loader = DataLoader(train_data, batch_size=96, num_workers=6) # Use the original training data
.. GENERATED FROM PYTHON SOURCE LINES 429-430
The validation data loader can be reused.
.. GENERATED FROM PYTHON SOURCE LINES 431-434
.. code-block:: default
valid_loader
.. rst-class:: sphx-glr-script-out
.. code-block:: none
<torch.utils.data.dataloader.DataLoader object at 0x7f5e187c0430>
.. GENERATED FROM PYTHON SOURCE LINES 435-438
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 439-455
.. 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
/data/data0/jiahang/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:447: 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.
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [3]
| Name | Type | Params
-----------------------------------------------
0 | criterion | CrossEntropyLoss | 0
1 | metrics | ModuleDict | 0
2 | model | DARTS | 345 K
-----------------------------------------------
345 K Trainable params
0 Non-trainable params
345 K Total params
1.381 Total estimated model params size (MB)
/data/data0/jiahang/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1892: 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.
"\n# Hello, NAS!\n\nThis is the 101 tutorial of Neural Architecture Search (NAS) on NNI.\nIn this tutorial, we will search for a neural architecture on MNIST dataset with the help of NAS framework of NNI, i.e., *Retiarii*.\nWe use multi-trial NAS as an example to show how to construct and explore a model space.\n\nThere are mainly three crucial components for a neural architecture search task, namely,\n\n* Model search space that defines a set of models to explore.\n* A proper strategy as the method to explore this model space.\n* A model evaluator that reports the performance of every model in the space.\n\nCurrently, PyTorch is the only supported framework by Retiarii, and we have only tested **PyTorch 1.7 to 1.10**.\nThis tutorial assumes PyTorch context but it should also apply to other frameworks, which is in our future plan.\n\n## Define your Model Space\n\nModel space is defined by users to express a set of models that users want to explore, which contains potentially good-performing models.\nIn this framework, a model space is defined with two parts: a base model and possible mutations on the base model.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Define Base Model\n\nDefining a base model is almost the same as defining a PyTorch (or TensorFlow) model.\nUsually, you only need to replace the code ``import torch.nn as nn`` with\n``import nni.retiarii.nn.pytorch as nn`` to use our wrapped PyTorch modules.\n\nBelow is a very simple example of defining a base model.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import torch\nimport torch.nn.functional as F\nimport nni.retiarii.nn.pytorch as nn\nfrom nni.retiarii import model_wrapper\n\n\n@model_wrapper # this decorator should be put on the out most\nclass Net(nn.Module):\n def __init__(self):\n super().__init__()\n self.conv1 = nn.Conv2d(1, 32, 3, 1)\n self.conv2 = nn.Conv2d(32, 64, 3, 1)\n self.dropout1 = nn.Dropout(0.25)\n self.dropout2 = nn.Dropout(0.5)\n self.fc1 = nn.Linear(9216, 128)\n self.fc2 = nn.Linear(128, 10)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n x = F.max_pool2d(self.conv2(x), 2)\n x = torch.flatten(self.dropout1(x), 1)\n x = self.fc2(self.dropout2(F.relu(self.fc1(x))))\n output = F.log_softmax(x, dim=1)\n return output"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
".. tip:: Always keep in mind that you should use ``import nni.retiarii.nn.pytorch as nn`` and :meth:`nni.retiarii.model_wrapper`.\n Many mistakes are a result of forgetting one of those.\n Also, please use ``torch.nn`` for submodules of ``nn.init``, e.g., ``torch.nn.init`` instead of ``nn.init``.\n\n### Define Model Mutations\n\nA base model is only one concrete model not a model space. We provide :doc:`API and Primitives </nas/construct_space>`\nfor users to express how the base model can be mutated. That is, to build a model space which includes many models.\n\nBased on the above base model, we can define a model space as below.\n\n.. code-block:: diff\n\n @model_wrapper\n class Net(nn.Module):\n def __init__(self):\n super().__init__()\n self.conv1 = nn.Conv2d(1, 32, 3, 1)\n - self.conv2 = nn.Conv2d(32, 64, 3, 1)\n + self.conv2 = nn.LayerChoice([\n + nn.Conv2d(32, 64, 3, 1),\n + DepthwiseSeparableConv(32, 64)\n + ])\n - self.dropout1 = nn.Dropout(0.25)\n + self.dropout1 = nn.Dropout(nn.ValueChoice([0.25, 0.5, 0.75]))\n self.dropout2 = nn.Dropout(0.5)\n - self.fc1 = nn.Linear(9216, 128)\n - self.fc2 = nn.Linear(128, 10)\n + feature = nn.ValueChoice([64, 128, 256])\n + self.fc1 = nn.Linear(9216, feature)\n + self.fc2 = nn.Linear(feature, 10)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n x = F.max_pool2d(self.conv2(x), 2)\n x = torch.flatten(self.dropout1(x), 1)\n x = self.fc2(self.dropout2(F.relu(self.fc1(x))))\n output = F.log_softmax(x, dim=1)\n return output\n\nThis results in the following code:\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"class DepthwiseSeparableConv(nn.Module):\n def __init__(self, in_ch, out_ch):\n super().__init__()\n self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size=3, groups=in_ch)\n self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1)\n\n def forward(self, x):\n return self.pointwise(self.depthwise(x))\n\n\n@model_wrapper\nclass ModelSpace(nn.Module):\n def __init__(self):\n super().__init__()\n self.conv1 = nn.Conv2d(1, 32, 3, 1)\n # LayerChoice is used to select a layer between Conv2d and DwConv.\n self.conv2 = nn.LayerChoice([\n nn.Conv2d(32, 64, 3, 1),\n DepthwiseSeparableConv(32, 64)\n ])\n # ValueChoice is used to select a dropout rate.\n # ValueChoice can be used as parameter of modules wrapped in `nni.retiarii.nn.pytorch`\n # or customized modules wrapped with `@basic_unit`.\n self.dropout1 = nn.Dropout(nn.ValueChoice([0.25, 0.5, 0.75])) # choose dropout rate from 0.25, 0.5 and 0.75\n self.dropout2 = nn.Dropout(0.5)\n feature = nn.ValueChoice([64, 128, 256])\n self.fc1 = nn.Linear(9216, feature)\n self.fc2 = nn.Linear(feature, 10)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n x = F.max_pool2d(self.conv2(x), 2)\n x = torch.flatten(self.dropout1(x), 1)\n x = self.fc2(self.dropout2(F.relu(self.fc1(x))))\n output = F.log_softmax(x, dim=1)\n return output\n\n\nmodel_space = ModelSpace()\nmodel_space"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This example uses two mutation APIs,\n:class:`nn.LayerChoice <nni.retiarii.nn.pytorch.LayerChoice>` and\n:class:`nn.InputChoice <nni.retiarii.nn.pytorch.ValueChoice>`.\n:class:`nn.LayerChoice <nni.retiarii.nn.pytorch.LayerChoice>`\ntakes a list of candidate modules (two in this example), one will be chosen for each sampled model.\nIt can be used like normal PyTorch module.\n:class:`nn.InputChoice <nni.retiarii.nn.pytorch.ValueChoice>` takes a list of candidate values,\none will be chosen to take effect for each sampled model.\n\nMore detailed API description and usage can be found :doc:`here </nas/construct_space>`.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>We are actively enriching the mutation APIs, to facilitate easy construction of model space.\n If the currently supported mutation APIs cannot express your model space,\n please refer to :doc:`this doc </nas/mutator>` for customizing mutators.</p></div>\n\n## Explore the Defined Model Space\n\nThere are basically two exploration approaches: (1) search by evaluating each sampled model independently,\nwhich is the search approach in `multi-trial NAS <multi-trial-nas>`\nand (2) one-shot weight-sharing based search, which is used in one-shot NAS.\nWe demonstrate the first approach in this tutorial. Users can refer to `here <one-shot-nas>` for the second approach.\n\nFirst, users need to pick a proper exploration strategy to explore the defined model space.\nSecond, users need to pick or customize a model evaluator to evaluate the performance of each explored model.\n\n### Pick an exploration strategy\n\nRetiarii supports many :doc:`exploration strategies </nas/exploration_strategy>`.\n\nSimply choosing (i.e., instantiate) an exploration strategy as below.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import nni.retiarii.strategy as strategy\nsearch_strategy = strategy.Random(dedup=True) # dedup=False if deduplication is not wanted"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Pick or customize a model evaluator\n\nIn the exploration process, the exploration strategy repeatedly generates new models. A model evaluator is for training\nand validating each generated model to obtain the model's performance.\nThe performance is sent to the exploration strategy for the strategy to generate better models.\n\nRetiarii has provided :doc:`built-in model evaluators </nas/evaluator>`, but to start with,\nit is recommended to use :class:`FunctionalEvaluator <nni.retiarii.evaluator.FunctionalEvaluator>`,\nthat is, to wrap your own training and evaluation code with one single function.\nThis function should receive one single model class and uses :func:`nni.report_final_result` to report the final score of this model.\n\nAn example here creates a simple evaluator that runs on MNIST dataset, trains for 2 epochs, and reports its validation accuracy.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import nni\n\nfrom torchvision import transforms\nfrom torchvision.datasets import MNIST\nfrom torch.utils.data import DataLoader\n\n\ndef train_epoch(model, device, train_loader, optimizer, epoch):\n loss_fn = torch.nn.CrossEntropyLoss()\n model.train()\n for batch_idx, (data, target) in enumerate(train_loader):\n data, target = data.to(device), target.to(device)\n optimizer.zero_grad()\n output = model(data)\n loss = loss_fn(output, target)\n loss.backward()\n optimizer.step()\n if batch_idx % 10 == 0:\n print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n epoch, batch_idx * len(data), len(train_loader.dataset),\n 100. * batch_idx / len(train_loader), loss.item()))\n\n\ndef test_epoch(model, device, test_loader):\n model.eval()\n test_loss = 0\n correct = 0\n with torch.no_grad():\n for data, target in test_loader:\n data, target = data.to(device), target.to(device)\n output = model(data)\n pred = output.argmax(dim=1, keepdim=True)\n correct += pred.eq(target.view_as(pred)).sum().item()\n\n test_loss /= len(test_loader.dataset)\n accuracy = 100. * correct / len(test_loader.dataset)\n\n print('\\nTest set: Accuracy: {}/{} ({:.0f}%)\\n'.format(\n correct, len(test_loader.dataset), accuracy))\n\n return accuracy\n\n\ndef evaluate_model(model_cls):\n # \"model_cls\" is a class, need to instantiate\n model = model_cls()\n\n device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n model.to(device)\n\n optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n transf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n train_loader = DataLoader(MNIST('data/mnist', download=True, transform=transf), batch_size=64, shuffle=True)\n test_loader = DataLoader(MNIST('data/mnist', download=True, train=False, transform=transf), batch_size=64)\n\n for epoch in range(3):\n # train the model for one epoch\n train_epoch(model, device, train_loader, optimizer, epoch)\n # test the model for one epoch\n accuracy = test_epoch(model, device, test_loader)\n # call report intermediate result. Result can be float or dict\n nni.report_intermediate_result(accuracy)\n\n # report final test result\n nni.report_final_result(accuracy)"
"The ``train_epoch`` and ``test_epoch`` here can be any customized function,\nwhere users can write their own training recipe.\n\nIt is recommended that the ``evaluate_model`` here accepts no additional arguments other than ``model_cls``.\nHowever, in the :doc:`advanced tutorial </nas/evaluator>`, we will show how to use additional arguments in case you actually need those.\nIn future, we will support mutation on the arguments of evaluators, which is commonly called \"Hyper-parmeter tuning\".\n\n## Launch an Experiment\n\nAfter all the above are prepared, it is time to start an experiment to do the model search. An example is shown below.\n\n"
"The following configurations are useful to control how many trials to run at most / at the same time.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"exp_config.max_trial_number = 4 # spawn 4 trials at most\nexp_config.trial_concurrency = 2 # will run two trials concurrently"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Remember to set the following config if you want to GPU.\n``use_active_gpu`` should be set true if you wish to use an occupied GPU (possibly running a GUI).\n\n"
"Launch the experiment. The experiment should take several minutes to finish on a workstation with 2 GPUs.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"exp.run(exp_config, 8081)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Users can also run Retiarii Experiment with :doc:`different training services </experiment/training_service/overview>`\nbesides ``local`` training service.\n\n## Visualize the Experiment\n\nUsers can visualize their experiment in the same way as visualizing a normal hyper-parameter tuning experiment.\nFor example, open ``localhost:8081`` in your browser, 8081 is the port that you set in ``exp.run``.\nPlease refer to :doc:`here </experiment/web_portal/web_portal>` for details.\n\nWe support visualizing models with 3rd-party visualization engines (like `Netron <https://netron.app/>`__).\nThis can be used by clicking ``Visualization`` in detail panel for each trial.\nNote that current visualization is based on `onnx <https://onnx.ai/>`__ ,\nthus visualization is not feasible if the model cannot be exported into onnx.\n\nBuilt-in evaluators (e.g., Classification) will automatically export the model into a file.\nFor your own evaluator, you need to save your file into ``$NNI_OUTPUT_DIR/model.onnx`` to make this work.\nFor instance,\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import os\nfrom pathlib import Path\n\n\ndef evaluate_model_with_visualization(model_cls):\n model = model_cls()\n # dump the model into an onnx\n if 'NNI_OUTPUT_DIR' in os.environ:\n dummy_input = torch.zeros(1, 3, 32, 32)\n torch.onnx.export(model, (dummy_input, ),\n Path(os.environ['NNI_OUTPUT_DIR']) / 'model.onnx')\n evaluate_model(model_cls)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Relaunch the experiment, and a button is shown on Web portal.\n\n<img src=\"file://../../img/netron_entrance_webui.png\">\n\n## Export Top Models\n\nUsers can export top models after the exploration is done using ``export_top_models``.\n\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"for model_dict in exp.export_top_models(formatter='dict'):\n print(model_dict)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The output is ``json`` object which records the mutation actions of the top model.\nIf users want to output source code of the top model,\nthey can use `graph-based execution engine <graph-based-execution-engine>` for the experiment,\nby simply adding the following two lines.\n\n"
"\n# Port PyTorch Quickstart to NNI\nThis is a modified version of `PyTorch quickstart`_.\n\nIt can be run directly and will have the exact same result as original version.\n\nFurthermore, it enables the ability of auto tuning with an NNI *experiment*, which will be detailed later.\n\nIt is recommended to run this script directly first to verify the environment.\n\nThere are 2 key differences from the original version:\n\n1. In `Get optimized hyperparameters`_ part, it receives generated hyperparameters.\n2. In `Train model and report accuracy`_ part, it reports accuracy metrics to NNI.\n\n"
"## Get optimized hyperparameters\nIf run directly, :func:`nni.get_next_parameter` is a no-op and returns an empty dict.\nBut with an NNI *experiment*, it will receive optimized hyperparameters from tuning algorithm.\n\n"
