备注

Click here to download the full example code

架构搜索入门教程¶

这是 NNI 上的神经架构搜索（NAS）的入门教程。在本教程中，我们将借助 NNI 的 NAS 框架，即 Retiarii，在 MNIST 数据集上实现网络结构搜索。我们以多尝试的架构搜索为例来展示如何构建和探索模型空间。

神经架构搜索任务主要有三个关键组成部分，即

模型搜索空间，定义了一个要探索的模型的集合。
一个合适的策略作为探索这个模型空间的方法。
一个模型评估器，用于为搜索空间中每个模型评估性能。

目前，Retiarii 只支持 PyTorch，并对 PyTorch 1.7 到 1.10 进行了测试。所以本教程假定您使用 PyTorch 作为深度学习框架。未来我们会支持更多框架。

定义您的模型空间¶

模型空间是由用户定义的，用来表达用户想要探索的一组模型，其中包含有潜力的好模型。在 NNI 的框架中，模型空间由两部分定义：基本模型和基本模型上可能的变化。

定义基本模型¶

定义基本模型与定义 PyTorch（或 TensorFlow）模型几乎相同。通常，您只需将代码 import torch.nn as nn 替换为 import nni.retiarii.nn.pytorch as nn 以使用我们打包的 PyTorch 模块。

下面是定义基本模型的一个非常简单的示例。

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

小技巧

记住，您应该使用 import nni.retiarii.nn.pytorch as nn 和 nni.retiarii.model_wrapper()。许多错误都是因为忘记使用某一个。另外，要使用 nn.init 的子模块，可以使用 torch.nn，例如， torch.nn.init 而不是 nn.init。

定义模型变化¶

基本模型只是一个具体模型，而不是模型空间。我们提供模型变化的 API 让用户表达如何改变基本模型。即构建一个包含许多模型的搜索空间。

基于上述基本模型，我们可以定义如下模型空间。

@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

结果是以下代码：

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

Out:

ModelSpace(
  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
  (conv2): LayerChoice([Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1)), DepthwiseSeparableConv(
    (depthwise): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), groups=32)
    (pointwise): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))
  )], label='model_1')
  (dropout1): Dropout(p=0.25, inplace=False)
  (dropout2): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=9216, out_features=64, bias=True)
  (fc2): Linear(in_features=64, out_features=10, bias=True)
)

这个例子使用了两个模型变化的 API， nn.LayerChoice 和 nn.InputChoice。 nn.LayerChoice 可以从一系列的候选子模块中（在本例中为两个），为每个采样模型选择一个。它可以像原来的 PyTorch 子模块一样使用。 nn.InputChoice 的参数是一个候选值列表，语义是为每个采样模型选择一个值。

更详细的 API 描述和用法可以在这里找到。

备注

我们正在积极丰富模型变化的 API，使得您可以轻松构建模型空间。如果当前支持的模型变化的 API 不能表达您的模型空间，请参考这篇文档来自定义突变。

探索定义的模型空间¶

简单来讲，有两种探索方法： (1) 独立评估每个采样到的模型，这是多尝试 NAS 中的搜索方法。 (2) 单尝试共享权重型的搜索，简称单尝试 NAS。我们在本教程中演示了第一种方法。第二种方法用户可以参考这里。

首先，用户需要选择合适的探索策略来探索定义好的模型空间。其次，用户需要选择或自定义模型性能评估来评估每个探索模型的性能。

选择探索策略¶

Retiarii 支持许多探索策略。

只需选择（即实例化）探索策略，就如下面的代码演示的一样：

import nni.retiarii.strategy as strategy
search_strategy = strategy.Random(dedup=True)  # dedup=False if deduplication is not wanted

Out:

/home/yugzhan/miniconda3/envs/cu102/lib/python3.8/site-packages/ray/autoscaler/_private/cli_logger.py:57: FutureWarning: Not all Ray CLI dependencies were found. In Ray 1.4+, the Ray CLI, autoscaler, and dashboard will only be usable via `pip install 'ray[default]'`. Please update your install command.
  warnings.warn(

挑选或自定义模型评估器¶

在探索过程中，探索策略反复生成新模型。模型评估器负责训练并验证每个生成的模型以获得模型的性能。该性能作为模型的得分被发送到探索策略以帮助其生成更好的模型。

Retiarii 提供了内置模型评估器，但在此之前，我们建议使用 FunctionalEvaluator，即用一个函数包装您自己的训练和评估代码。这个函数应该接收一个单一的模型类并使用 nni.report_final_result() 报告这个模型的最终分数。

此处的示例创建了一个简单的评估器，该评估器在 MNIST 数据集上运行，训练 2 个 epoch，并报告其在验证集上的准确率。

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)

创建评估器

from nni.retiarii.evaluator import FunctionalEvaluator
evaluator = FunctionalEvaluator(evaluate_model)

这里的 train_epoch 和 test_epoch 可以是任何自定义函数，用户可以在其中编写自己的训练逻辑。

建议这里的 evaluate_model 不接受除 model_cls 之外的其他参数。但是，在高级教程 </nas/evaluator> 中，我们将展示如何使用其他参数，以免您确实需要这些参数。未来，我们将支持对评估器的参数进行变化（通常称为“超参数调优”）。

启动实验¶

一切都已准备就绪，现在就可以开始做模型搜索的实验了。如下所示。

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'

以下配置可以用于控制最多/同时运行多少试验。

exp_config.max_trial_number = 4   # 最多运行 4 个实验
exp_config.trial_concurrency = 2  # 最多同时运行 2 个试验

如果要使用 GPU，请设置以下配置。如果您希望使用被占用了的 GPU（比如 GPU 上可能正在运行 GUI），则 use_active_gpu 应设置为 true。

exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = True

启动实验。在一个有两块 GPU 的工作站上完成整个实验大约需要几分钟时间。

exp.run(exp_config, 8081)

Out:

INFO:nni.experiment:Creating experiment, Experiment ID: z8ns5fv7
INFO:nni.experiment:Connecting IPC pipe...
INFO:nni.experiment:Starting web server...
INFO:nni.experiment:Setting up...
INFO:nni.runtime.msg_dispatcher_base:Dispatcher started
INFO:nni.retiarii.experiment.pytorch:Web UI URLs: http://127.0.0.1:8081 http://10.190.172.35:8081 http://192.168.49.1:8081 http://172.17.0.1:8081
INFO:nni.retiarii.experiment.pytorch:Start strategy...
INFO:root:Successfully update searchSpace.
INFO:nni.retiarii.strategy.bruteforce:Random search running in fixed size mode. Dedup: on.
INFO:nni.retiarii.experiment.pytorch:Stopping experiment, please wait...
INFO:nni.retiarii.experiment.pytorch:Strategy exit
INFO:nni.retiarii.experiment.pytorch:Waiting for experiment to become DONE (you can ctrl+c if there is no running trial jobs)...
INFO:nni.runtime.msg_dispatcher_base:Dispatcher exiting...
INFO:nni.retiarii.experiment.pytorch:Experiment stopped

除了 local 训练平台，用户还可以使用不同的训练平台来运行 Retiarii 试验。

可视化实验¶

用户可以可视化他们的架构搜索实验，就像可视化超参调优实验一样。例如，在浏览器中打开 localhost:8081，8081 是您在 exp.run 中设置的端口。详情请参考这里。

我们支持使用第三方可视化引擎（如 Netron）对模型进行可视化。这可以通过单击每个试验的详细面板中的“可视化”来使用。请注意，当前的可视化是基于 onnx，因此，如果模型不能导出为 onnx，可视化是不可行的。

内置评估器（例如 Classification）会将模型自动导出到文件中。对于您自己的评估器，您需要将文件保存到 $NNI_OUTPUT_DIR/model.onnx。例如，

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)

重新启动实验，Web 界面上会显示一个按钮。

导出最优模型¶

搜索完成后，用户可以使用 export_top_models 导出最优模型。

for model_dict in exp.export_top_models(formatter='dict'):
    print(model_dict)

Out:

{'model_1': '0', 'model_2': 0.25, 'model_3': 64}

输出是一个 JSON 对象，记录了最好的模型的每一个选择都选了什么。如果用户想要搜出来的模型的源代码，他们可以使用基于图的引擎，只需增加如下两行。

exp_config.execution_engine = 'base'
export_formatter = 'code'

Total running time of the script: ( 2 minutes 4.499 seconds)

Gallery generated by Sphinx-Gallery