NAS tutorials (#4573)

docs/source/tutorials/sg_execution_times.rst

@@ -5,10 +5,12 @@
 
 Computation times
 =================
-**00:24.663** total execution time for **tutorials** files:
+**02:24.722** total execution time for **tutorials** files:
 
-+-----------------------------------------------------------------------------------+-----------+--------+
-| :ref:`sphx_glr_tutorials_nni_experiment.py` (``nni_experiment.py``)               | 00:24.662 | 0.0 MB |
-+-----------------------------------------------------------------------------------+-----------+--------+
-| :ref:`sphx_glr_tutorials_nas_quick_start_mnist.py` (``nas_quick_start_mnist.py``) | 00:00.002 | 0.0 MB |
-+-----------------------------------------------------------------------------------+-----------+--------+
++-------------------------------------------------------------------------------+-----------+--------+
+| :ref:`sphx_glr_tutorials_hello_nas.py` (``hello_nas.py``)                     | 02:24.722 | 0.0 MB |
++-------------------------------------------------------------------------------+-----------+--------+
+| :ref:`sphx_glr_tutorials_nasbench_as_dataset.py` (``nasbench_as_dataset.py``) | 00:00.000 | 0.0 MB |
++-------------------------------------------------------------------------------+-----------+--------+
+| :ref:`sphx_glr_tutorials_nni_experiment.py` (``nni_experiment.py``)           | 00:00.000 | 0.0 MB |
++-------------------------------------------------------------------------------+-----------+--------+

examples/tutorials/.gitignore

+data/

examples/tutorials/hello_nas.py

+"""
+Hello, NAS!
+===========
+
+This is the 101 tutorial of Neural Architecture Search (NAS) on NNI.
+In this tutorial, we will search for a neural architecture on MNIST dataset with the help of NAS framework of NNI, i.e., *Retiarii*.
+We use multi-trial NAS as an example to show how to construct and explore a model space.
+
+There are mainly three crucial components for a neural architecture search task, namely,
+
+* Model search space that defines a set of models to explore.
+* A proper strategy as the method to explore this model space.
+* A model evaluator that reports the performance of every model in the space.
+
+Currently, PyTorch is the only supported framework by Retiarii, and we have only tested **PyTorch 1.7 to 1.10**.
+This tutorial assumes PyTorch context but it should also apply to other frameworks, which is in our future plan.
+
+Define your Model Space
+-----------------------
+
+Model space is defined by users to express a set of models that users want to explore, which contains potentially good-performing models.
+In this framework, a model space is defined with two parts: a base model and possible mutations on the base model.
+"""
+
+# %%
+#
+# Define Base Model
+# ^^^^^^^^^^^^^^^^^
+#
+# Defining a base model is almost the same as defining a PyTorch (or TensorFlow) model.
+# Usually, you only need to replace the code ``import torch.nn as nn`` with
+# ``import nni.retiarii.nn.pytorch as nn`` to use our wrapped PyTorch modules.
+#
+# Below is a very simple example of defining a base model.
+
+import torch
+import torch.nn.functional as F
+import nni.retiarii.nn.pytorch as nn
+from nni.retiarii import model_wrapper
+
+
+@model_wrapper      # this decorator should be put on the out most
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.max_pool2d(self.conv2(x), 2)
+        x = torch.flatten(self.dropout1(x), 1)
+        x = self.fc2(self.dropout2(F.relu(self.fc1(x))))
+        output = F.log_softmax(x, dim=1)
+        return output
+
+# %%
+# .. tip:: Always keep in mind that you should use ``import nni.retiarii.nn.pytorch as nn`` and :meth:`nni.retiarii.model_wrapper`.
+#          Many mistakes are a result of forgetting one of those.
+#          Also, please use ``torch.nn`` for submodules of ``nn.init``, e.g., ``torch.nn.init`` instead of ``nn.init``.
+#
+# Define Model Mutations
+# ^^^^^^^^^^^^^^^^^^^^^^
+#
+# A base model is only one concrete model not a model space. We provide :doc:`API and Primitives </NAS/MutationPrimitives>`
+# for users to express how the base model can be mutated. That is, to build a model space which includes many models.
+#
+# Based on the above base model, we can define a model space as below.
+#
+# .. code-block:: diff
+#
+#   @model_wrapper
+#   class Net(nn.Module):
+#     def __init__(self):
+#       super().__init__()
+#       self.conv1 = nn.Conv2d(1, 32, 3, 1)
+#   -   self.conv2 = nn.Conv2d(32, 64, 3, 1)
+#   +   self.conv2 = nn.LayerChoice([
+#   +       nn.Conv2d(32, 64, 3, 1),
+#   +       DepthwiseSeparableConv(32, 64)
+#   +   ])
+#   -   self.dropout1 = nn.Dropout(0.25)
+#   +   self.dropout1 = nn.Dropout(nn.ValueChoice([0.25, 0.5, 0.75]))
+#       self.dropout2 = nn.Dropout(0.5)
+#   -   self.fc1 = nn.Linear(9216, 128)
+#   -   self.fc2 = nn.Linear(128, 10)
+#   +   feature = nn.ValueChoice([64, 128, 256])
+#   +   self.fc1 = nn.Linear(9216, feature)
+#   +   self.fc2 = nn.Linear(feature, 10)
+#
+#     def forward(self, x):
+#       x = F.relu(self.conv1(x))
+#       x = F.max_pool2d(self.conv2(x), 2)
+#       x = torch.flatten(self.dropout1(x), 1)
+#       x = self.fc2(self.dropout2(F.relu(self.fc1(x))))
+#       output = F.log_softmax(x, dim=1)
+#       return output
+#
+# This results in the following code:
+
+
+class DepthwiseSeparableConv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size=3, groups=in_ch)
+        self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1)
+
+    def forward(self, x):
+        return self.pointwise(self.depthwise(x))
+
+
+@model_wrapper
+class ModelSpace(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        # LayerChoice is used to select a layer between Conv2d and DwConv.
+        self.conv2 = nn.LayerChoice([
+            nn.Conv2d(32, 64, 3, 1),
+            DepthwiseSeparableConv(32, 64)
+        ])
+        # ValueChoice is used to select a dropout rate.
+        # ValueChoice can be used as parameter of modules wrapped in `nni.retiarii.nn.pytorch`
+        # or customized modules wrapped with `@basic_unit`.
+        self.dropout1 = nn.Dropout(nn.ValueChoice([0.25, 0.5, 0.75]))  # choose dropout rate from 0.25, 0.5 and 0.75
+        self.dropout2 = nn.Dropout(0.5)
+        feature = nn.ValueChoice([64, 128, 256])
+        self.fc1 = nn.Linear(9216, feature)
+        self.fc2 = nn.Linear(feature, 10)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.max_pool2d(self.conv2(x), 2)
+        x = torch.flatten(self.dropout1(x), 1)
+        x = self.fc2(self.dropout2(F.relu(self.fc1(x))))
+        output = F.log_softmax(x, dim=1)
+        return output
+
+
+model_space = ModelSpace()
+model_space
+
+# %%
+# This example uses two mutation APIs, ``nn.LayerChoice`` and ``nn.ValueChoice``.
+# ``nn.LayerChoice`` takes a list of candidate modules (two in this example), one will be chosen for each sampled model.
+# It can be used like normal PyTorch module.
+# ``nn.ValueChoice`` takes a list of candidate values, one will be chosen to take effect for each sampled model.
+#
+# More detailed API description and usage can be found :doc:`here </NAS/construct_space>`.
+#
+# .. note::
+#
+#     We are actively enriching the mutation APIs, to facilitate easy construction of model space.
+#     If the currently supported mutation APIs cannot express your model space,
+#     please refer to :doc:`this doc </NAS/Mutators>` for customizing mutators.
+#
+# Explore the Defined Model Space
+# -------------------------------
+#
+# There are basically two exploration approaches: (1) search by evaluating each sampled model independently,
+# which is the search approach in multi-trial NAS and (2) one-shot weight-sharing based search, which is used in one-shot NAS.
+# We demonstrate the first approach in this tutorial. Users can refer to :doc:`here </NAS/OneshotTrainer>` for the second approach.
+#
+# First, users need to pick a proper exploration strategy to explore the defined model space.
+# Second, users need to pick or customize a model evaluator to evaluate the performance of each explored model.
+#
+# Pick an exploration strategy
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+#
+# Retiarii supports many :doc:`exploration strategies </NAS/ExplorationStrategies>`.
+#
+# Simply choosing (i.e., instantiate) an exploration strategy as below.
+
+import nni.retiarii.strategy as strategy
+search_strategy = strategy.Random(dedup=True)  # dedup=False if deduplication is not wanted
+
+# %%
+# Pick or customize a model evaluator
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+#
+# In the exploration process, the exploration strategy repeatedly generates new models. A model evaluator is for training and validating each generated model to obtain the model's performance. The performance is sent to the exploration strategy for the strategy to generate better models.
+#
+# Retiarii has provided :doc:`built-in model evaluators </NAS/ModelEvaluators>`, but to start with, it is recommended to use ``FunctionalEvaluator``, that is, to wrap your own training and evaluation code with one single function. This function should receive one single model class and uses ``nni.report_final_result`` to report the final score of this model.
+#
+# An example here creates a simple evaluator that runs on MNIST dataset, trains for 2 epochs, and reports its validation accuracy.
+
+import nni
+
+from torchvision import transforms
+from torchvision.datasets import MNIST
+from torch.utils.data import DataLoader
+
+
+def train_epoch(model, device, train_loader, optimizer, epoch):
+    loss_fn = torch.nn.CrossEntropyLoss()
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = loss_fn(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % 10 == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss.item()))
+
+
+def test_epoch(model, device, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+    accuracy = 100. * correct / len(test_loader.dataset)
+
+    print('\nTest set: Accuracy: {}/{} ({:.0f}%)\n'.format(
+          correct, len(test_loader.dataset), accuracy))
+
+    return accuracy
+
+
+def evaluate_model(model_cls):
+    # "model_cls" is a class, need to instantiate
+    model = model_cls()
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    model.to(device)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    transf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
+    train_loader = DataLoader(MNIST('data/mnist', download=True, transform=transf), batch_size=64, shuffle=True)
+    test_loader = DataLoader(MNIST('data/mnist', download=True, train=False, transform=transf), batch_size=64)
+
+    for epoch in range(3):
+        # train the model for one epoch
+        train_epoch(model, device, train_loader, optimizer, epoch)
+        # test the model for one epoch
+        accuracy = test_epoch(model, device, test_loader)
+        # call report intermediate result. Result can be float or dict
+        nni.report_intermediate_result(accuracy)
+
+    # report final test result
+    nni.report_final_result(accuracy)
+
+
+# %%
+# Create the evaluator
+
+from nni.retiarii.evaluator import FunctionalEvaluator
+evaluator = FunctionalEvaluator(evaluate_model)
+
+# %%
+#
+# The ``train_epoch`` and ``test_epoch`` here can be any customized function, where users can write their own training recipe.
+#
+# It is recommended that the :doc:``evaluate_model`` here accepts no additional arguments other than ``model_cls``.
+# However, in the `advanced tutorial </NAS/ModelEvaluators>`, we will show how to use additional arguments in case you actually need those.
+# In future, we will support mutation on the arguments of evaluators, which is commonly called "Hyper-parmeter tuning".
+#
+# Launch an Experiment
+# --------------------
+#
+# After all the above are prepared, it is time to start an experiment to do the model search. An example is shown below.
+
+from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
+exp = RetiariiExperiment(model_space, evaluator, [], search_strategy)
+exp_config = RetiariiExeConfig('local')
+exp_config.experiment_name = 'mnist_search'
+
+# %%
+# The following configurations are useful to control how many trials to run at most / at the same time.
+
+exp_config.max_trial_number = 4   # spawn 4 trials at most
+exp_config.trial_concurrency = 2  # will run two trials concurrently
+
+# %%
+# Remember to set the following config if you want to GPU.
+# ``use_active_gpu`` should be set true if you wish to use an occupied GPU (possibly running a GUI).
+
+exp_config.trial_gpu_number = 1
+exp_config.training_service.use_active_gpu = False
+
+# %%
+# Launch the experiment. The experiment should take several minutes to finish on a workstation with 2 GPUs.
+
+exp.run(exp_config, 8081)
+
+# %%
+# Users can also run Retiarii Experiment with :doc:`different training services <../training_services>` besides ``local`` training service.
+#
+# Visualize the Experiment
+# ------------------------
+#
+# Users can visualize their experiment in the same way as visualizing a normal hyper-parameter tuning experiment.
+# For example, open ``localhost:8081`` in your browser, 8081 is the port that you set in ``exp.run``.
+# Please refer to :doc:`here <../Tutorial/WebUI>` for details.
+#
+# We support visualizing models with 3rd-party visualization engines (like `Netron <https://netron.app/>`__).
+# This can be used by clicking ``Visualization`` in detail panel for each trial.
+# Note that current visualization is based on `onnx <https://onnx.ai/>`__ ,
+# thus visualization is not feasible if the model cannot be exported into onnx.
+#
+# Built-in evaluators (e.g., Classification) will automatically export the model into a file.
+# For your own evaluator, you need to save your file into ``$NNI_OUTPUT_DIR/model.onnx`` to make this work.
+# For instance,
+
+import os
+from pathlib import Path
+
+
+def evaluate_model_with_visualization(model_cls):
+    model = model_cls()
+    # dump the model into an onnx
+    if 'NNI_OUTPUT_DIR' in os.environ:
+        dummy_input = torch.zeros(1, 3, 32, 32)
+        torch.onnx.export(model, (dummy_input, ),
+                          Path(os.environ['NNI_OUTPUT_DIR']) / 'model.onnx')
+    evaluate_model(model_cls)
+
+# %%
+# Relaunch the experiment, and a button is shown on WebUI.
+#
+# .. image:: ../../img/netron_entrance_webui.png
+#
+# Export Top Models
+# -----------------
+#
+# Users can export top models after the exploration is done using ``export_top_models``.
+
+for model_dict in exp.export_top_models(formatter='dict'):
+    print(model_dict)
+
+# The output is `json` object which records the mutation actions of the top model.
+# If users want to output source code of the top model, they can use graph-based execution engine for the experiment,
+# by simply adding the following two lines.
+#
+# .. code-block:: python
+#
+#   exp_config.execution_engine = 'base'
+#   export_formatter = 'code'

examples/tutorials/nas_quick_start_mnist.py

-"""
-Get started with NAS on MNIST
-=============================
-"""
-
-# %%
-a = (1, 2, 3)
-a
-
-# %%
-print('hello')

examples/tutorials/nasbench_as_dataset.py

+"""
+Use NAS Benchmarks as Datasets
+==============================
+
+In this tutorial, we show how to use NAS Benchmarks as datasets.
+For research purposes we sometimes desire to query the benchmarks for architecture accuracies,
+rather than train them one by one from scratch.
+NNI has provided query tools so that users can easily get the retrieve the data in NAS benchmarks.
+"""
+
+# %%
+# Prerequisites
+# -------------
+# This tutorial assumes that you have already prepared your NAS benchmarks under cache directory
+# (by default, ``~/.cache/nni/nasbenchmark``).
+# If you haven't, please follow the data preparation guide in :doc:`../NAS/Benchmarks`.
+#
+# As a result, the directory should look like:
+
+import os
+os.listdir(os.path.expanduser('~/.cache/nni/nasbenchmark'))
+
+# %%
+import pprint
+
+from nni.nas.benchmarks.nasbench101 import query_nb101_trial_stats
+from nni.nas.benchmarks.nasbench201 import query_nb201_trial_stats
+from nni.nas.benchmarks.nds import query_nds_trial_stats
+
+# %%
+# NAS-Bench-101
+# -------------
+#
+# Use the following architecture as an example:
+#
+# .. image:: ../../img/nas-bench-101-example.png
+
+arch = {
+    'op1': 'conv3x3-bn-relu',
+    'op2': 'maxpool3x3',
+    'op3': 'conv3x3-bn-relu',
+    'op4': 'conv3x3-bn-relu',
+    'op5': 'conv1x1-bn-relu',
+    'input1': [0],
+    'input2': [1],
+    'input3': [2],
+    'input4': [0],
+    'input5': [0, 3, 4],
+    'input6': [2, 5]
+}
+for t in query_nb101_trial_stats(arch, 108, include_intermediates=True):
+    pprint.pprint(t)
+
+# %%
+# An architecture of NAS-Bench-101 could be trained more than once.
+# Each element of the returned generator is a dict which contains one of the training results of this trial config
+# (architecture + hyper-parameters) including train/valid/test accuracy,
+# training time, number of epochs, etc. The results of NAS-Bench-201 and NDS follow similar formats.
+#
+# NAS-Bench-201
+# -------------
+#
+# Use the following architecture as an example:
+#
+# .. image:: ../../img/nas-bench-201-example.png
+
+arch = {
+    '0_1': 'avg_pool_3x3',
+    '0_2': 'conv_1x1',
+    '1_2': 'skip_connect',
+    '0_3': 'conv_1x1',
+    '1_3': 'skip_connect',
+    '2_3': 'skip_connect'
+}
+for t in query_nb201_trial_stats(arch, 200, 'cifar100'):
+    pprint.pprint(t)
+
+# %%
+# Intermediate results are also available.
+
+for t in query_nb201_trial_stats(arch, None, 'imagenet16-120', include_intermediates=True):
+    print(t['config'])
+    print('Intermediates:', len(t['intermediates']))
+
+# %%
+# NDS
+# ---
+#
+# Use the following architecture as an example:
+#
+# .. image:: ../../img/nas-bench-nds-example.png
+#
+# Here, ``bot_muls``, ``ds``, ``num_gs``, ``ss`` and ``ws`` stand for "bottleneck multipliers",
+# "depths", "number of groups", "strides" and "widths" respectively.
+
+# %%
+model_spec = {
+    'bot_muls': [0.0, 0.25, 0.25, 0.25],
+    'ds': [1, 16, 1, 4],
+    'num_gs': [1, 2, 1, 2],
+    'ss': [1, 1, 2, 2],
+    'ws': [16, 64, 128, 16]
+}
+
+# %%
+# Use none as a wildcard.
+for t in query_nds_trial_stats('residual_bottleneck', None, None, model_spec, None, 'cifar10'):
+    pprint.pprint(t)
+
+# %%
+model_spec = {
+    'bot_muls': [0.0, 0.25, 0.25, 0.25],
+    'ds': [1, 16, 1, 4],
+    'num_gs': [1, 2, 1, 2],
+    'ss': [1, 1, 2, 2],
+    'ws': [16, 64, 128, 16]
+}
+for t in query_nds_trial_stats('residual_bottleneck', None, None, model_spec, None, 'cifar10', include_intermediates=True):
+    pprint.pprint(t['intermediates'][:10])
+
+# %%
+model_spec = {'ds': [1, 12, 12, 12], 'ss': [1, 1, 2, 2], 'ws': [16, 24, 24, 40]}
+for t in query_nds_trial_stats('residual_basic', 'resnet', 'random', model_spec, {}, 'cifar10'):
+    pprint.pprint(t)
+
+# %%
+# Get the first one.
+pprint.pprint(next(query_nds_trial_stats('vanilla', None, None, None, None, None)))
+
+# %%
+# Count number.
+model_spec = {'num_nodes_normal': 5, 'num_nodes_reduce': 5, 'depth': 12, 'width': 32, 'aux': False, 'drop_prob': 0.0}
+cell_spec = {
+    'normal_0_op_x': 'avg_pool_3x3',
+    'normal_0_input_x': 0,
+    'normal_0_op_y': 'conv_7x1_1x7',
+    'normal_0_input_y': 1,
+    'normal_1_op_x': 'sep_conv_3x3',
+    'normal_1_input_x': 2,
+    'normal_1_op_y': 'sep_conv_5x5',
+    'normal_1_input_y': 0,
+    'normal_2_op_x': 'dil_sep_conv_3x3',
+    'normal_2_input_x': 2,
+    'normal_2_op_y': 'dil_sep_conv_3x3',
+    'normal_2_input_y': 2,
+    'normal_3_op_x': 'skip_connect',
+    'normal_3_input_x': 4,
+    'normal_3_op_y': 'dil_sep_conv_3x3',
+    'normal_3_input_y': 4,
+    'normal_4_op_x': 'conv_7x1_1x7',
+    'normal_4_input_x': 2,
+    'normal_4_op_y': 'sep_conv_3x3',
+    'normal_4_input_y': 4,
+    'normal_concat': [3, 5, 6],
+    'reduce_0_op_x': 'avg_pool_3x3',
+    'reduce_0_input_x': 0,
+    'reduce_0_op_y': 'dil_sep_conv_3x3',
+    'reduce_0_input_y': 1,
+    'reduce_1_op_x': 'sep_conv_3x3',
+    'reduce_1_input_x': 0,
+    'reduce_1_op_y': 'sep_conv_3x3',
+    'reduce_1_input_y': 0,
+    'reduce_2_op_x': 'skip_connect',
+    'reduce_2_input_x': 2,
+    'reduce_2_op_y': 'sep_conv_7x7',
+    'reduce_2_input_y': 0,
+    'reduce_3_op_x': 'conv_7x1_1x7',
+    'reduce_3_input_x': 4,
+    'reduce_3_op_y': 'skip_connect',
+    'reduce_3_input_y': 4,
+    'reduce_4_op_x': 'conv_7x1_1x7',
+    'reduce_4_input_x': 0,
+    'reduce_4_op_y': 'conv_7x1_1x7',
+    'reduce_4_input_y': 5,
+    'reduce_concat': [3, 6]
+}
+
+for t in query_nds_trial_stats('nas_cell', None, None, model_spec, cell_spec, 'cifar10'):
+    assert t['config']['model_spec'] == model_spec
+    assert t['config']['cell_spec'] == cell_spec
+    pprint.pprint(t)
+
+# %%
+# Count number.
+print('NDS (amoeba) count:', len(list(query_nds_trial_stats(None, 'amoeba', None, None, None, None, None))))