DARTS搜索空间

在DARTS搜索空间中搜索

在本教程中，我们将演示如何在DARTS_中搜索著名的模型空间。

通过这个过程，您将学会：

如何使用NNI的模型空间中心提供的内置模型空间。
如何使用一次性的探索策略来探索模型空间。
如何自定义评估器以实现最佳性能。

最后，我们在CIFAR-10数据集上得到了一个表现强劲的模型，其准确率达到了97.28%。

.. 注意::

运行此教程需要一个GPU。
如果没有GPU，您可以在:class:~nni.nas.evaluator.pytorch.Classification中将gpus设置为0，
但请注意这会慢得多。

使用预搜索的DARTS模型

与PyTorch的入门教程_类似，
我们从CIFAR-10数据集开始，这是一个包含10个类别的图像分类数据集。
CIFAR-10中的图像大小为3x32x32，即32x32像素大小的RGB彩色图像。

我们首先使用torchvision加载CIFAR-10数据集。

import nni
import torch
from torchvision import transforms
from torchvision.datasets import CIFAR10
from nni.nas.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)

注意

如果您要使用多试验策略，则必须使用 :func:`nni.trace` 将CIFAR10包装起来，并使用来自``nni.nas.evaluator.pytorch``的DataLoader（而不是``torch.utils.data``）。否则，这是可选的。

NNI提供了许多内置模型空间，以及:doc:模型空间中心 </nas/space_hub>中的许多预搜索模型，
这些模型是由最流行的NAS文献产生的。
预训练模型是以前在诸如CIFAR-10或ImageNet等大型数据集上先前训练过的保存的网络。
您可以轻松地将这些模型加载为起点，验证其性能，并在需要时进行微调。

在本教程中，我们选择DARTS_搜索空间中的一个模型，该模型在目标数据集CIFAR-10上原生训练，
以节省繁琐的微调步骤。

.. 提示::

在其他数据集上微调预搜索模型与微调任何模型没有区别。
如果您想了解在PyTorch中通常如何进行微调，
我们建议阅读目标检测微调的本教程_。

from nni.nas.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, cuda=True)  # 如果没有GPU请将此设置为false。

使用预搜索模型的旅程就到此为止了。或者如果您有兴趣，我们可以进一步在:class:~nni.nas.hub.pytorch.DARTS空间中搜索模型。

使用DARTS模型空间

DARTS论文中提供的模型空间起源于NASNet，
其中完整的模型由重复堆叠单个计算单元（称为cell）构建。
网络中有两种类型的cell，一种称为normal cell，另一种称为reduction cell。
normal cell和reduction cell之间的关键区别在于reduction cell会对输入特征图进行下采样，
并降低其分辨率。normal cell和reduction cell交替堆叠，如下图所示。

cell接受两个前一个cell的输出作为输入，并包含一系列节点。
每个节点接受同一个cell中的两个前一个节点（或两个cell输入），
并对每个输入应用一个算子（例如卷积或最大池化），
并将算子的输出求和作为节点的输出。
cell的输出是所有从未被用作另一个节点的输入的节点的连接。
如果想了解更多细节，可以阅读NDS或ENAS。

下图演示了cells的一个示例。

DARTS_论文中提出的搜索空间对原始空间NASNet_进行了两个修改。

首先，操作符候选人已经缩小到七个：

最大池化3x3
平均池化3x3
跳跃连接（Identity）
可分离卷积3x3
可分离卷积5x5
空洞卷积3x3
空洞卷积5x5

其次，cell的输出是cell中的所有节点的连接。

由于搜索空间基于cell，一旦固定了normal和reduction cell，我们就可以无限次地堆叠它们。为了节省搜索成本，
在搜索阶段通常会减少滤波器（即通道）的数量和堆叠的cell的数量，并在训练最终搜索的架构时增加它们的数量。

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-1.md

注意

`DARTS`_ 是一篇在搜索空间和搜索策略上都有创新的论文。在本教程中，我们将使用 DARTS 提供的**模型空间**和 DARTS 提出的**搜索策略**进行搜索。我们将它们分别称为 *DARTS 模型空间*（``DartsSpace``）和 *DARTS 策略*（``DartsStrategy``）。我们并没有意味着需要同时使用 :class:`~nni.nas.hub.pytorch.DARTS` 空间和 :class:`~nni.nas.strategy.DARTS` 策略。您始终可以使用其他搜索策略探索 DARTS 空间，或使用自己的策略搜索不同的模型空间。

在下面的示例中，我们使用 16 个初始过滤器和 8 个堆叠单元初始化一个 :class:~nni.nas.hub.pytorch.DARTS 模型空间。
该网络专门用于 CIFAR-10 数据集，输入分辨率为 32x32。

这里的 :class:~nni.nas.hub.pytorch.DARTS 模型空间是由 :doc:模型空间中心 </nas/space_hub>提供的，
我们已经支持了多个流行的模型空间以供插拔使用。

.. tip::

这里的模型空间可以用中心提供的任何空间代替，
甚至可以使用从零开始构建的自定义空间。

model_space = DartsSpace(
    width=16,           # 模型的初始过滤器（通道数）
    num_cells=8,        # 总共堆叠的单元数
    dataset='cifar'     # 给出输入分辨率的提示，这里是 32x32
)

在模型空间上搜索

警告

请将 ``fast_dev_run`` 设置为 False 以重现我们声称的结果。否则，只会运行少量的小批次。

fast_dev_run = True

评估器

要开始探索模型空间，首先需要有一个评估器来提供“好模型”的标准。
由于我们在 CIFAR-10 数据集上进行搜索，您可以将 :class:~nni.nas.evaluator.pytorch.Classification 作为起点。

请注意，在 NAS 的典型设置中，模型搜索应在验证集上进行，最终搜索到的模型的评估应在测试集上进行。
但是，由于 CIFAR-10 数据集没有测试数据集（只有 50k 训练集 + 10k 验证集），
我们必须将原始训练集分成一个训练集和一个验证集。
Weight-Sharing Neural Architecture Search: A Battle to Shrink the Optimization Gap_
和
How Does Supernet Help in Neural Architecture Search?_。请感兴趣的读者参阅。

:class:~nni.nas.strategy.DARTS 策略是 NNI 提供的 :doc:内置搜索策略之一 </nas/exploration_strategy>。
使用它只需要一行代码。

from nni.nas.strategy import DARTS as DartsStrategy

strategy = DartsStrategy()

.. tip:: 这里的 DartsStrategy 可以替换为任何搜索策略，甚至是多次试验的策略。

如果您想了解 DARTS 策略的工作原理，这里有一个简要版本。
在幕后，DARTS 将单元转换为一个密集连接图，并在边上放置运算符（见下图）。
由于运算符尚未确定，每个边都是多个运算符的加权混合（图中有多种颜色）。
DARTS 然后学习为每个边分配最佳的“颜色”来进行网络训练。
最后，它选择每个边的一种“颜色”，并且删除多余的边。
边上的权重称为架构权重。

.. tip:: 图中没有反映出来的一点是，对于 DARTS 模型空间，每个节点保留了两个输入。

启动实验

然后我们来到启动实验的步骤。
这个步骤类似于我们在 :doc:入门教程 <hello_nas> 中所做的。

from nni.nas.experiment import NasExperiment

experiment = NasExperiment(model_space, evaluator, strategy)
experiment.run()

.. tip::

搜索过程可以用 tensorboard 可视化。例如::

   tensorboard --logdir=./lightning_logs

然后，打开浏览器并转到 http://localhost:6006/ 监视搜索过程。

.. image:: ../../img/darts_search_process.png

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-2.md

我们可以通过 export_top_models 方法检索策略找到的最佳模型。这里，检索到的模型是一个称为“architecture dict”的字典，描述了选定的正常单元和缩减单元。

exported_arch = experiment.export_top_models(formatter='dict')[0]

exported_arch

单元可以使用以下代码段可视化（从 DARTS visualization 中复制并修改）。

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)

警告

上面的单元是通过 ``fast_dev_run`` 获得的（即只运行1个 mini-batch）。

当关闭 fast_dev_run 时，我们将获得以下架构的模型，您可能注意到有趣的事实是，约有一半的操作选择了 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
})

重新训练搜索到的模型

上一步得到的只是一个单元结构。为了得到一个带有训练权重的最终可用模型，我们需要基于这个结构构建一个真实模型，然后进行全面训练。

为了基于从实验中导出的架构字典构建一个固定模型，我们可以使用 :func:nni.nas.space.model_context 方法。在上下文内部，我们将基于 exported_arch 创建一个固定模型，而不是创建一个空间。

from nni.nas.space import model_context

with model_context(exported_arch):
    final_model = DartsSpace(width=16, num_cells=8, dataset='cifar')

然后，我们在完整的 CIFAR-10 训练数据集上训练模型，并在原始的 CIFAR-10 验证数据集上进行评估。

train_loader = DataLoader(train_data, batch_size=96, num_workers=6)  # 使用原始的训练数据

验证数据加载器可以重复使用。

valid_loader

在此处我们必须创建一个新的 evaluator，因为使用了不同的数据拆分。此外，我们应该避免基于 :class:~nni.nas.evaluator.pytorch.Classification 加载错误检查点的底层 pytorch-lightning 实现。

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-3.md

max_epochs = 100

evaluator = 分类(
    学习率=1e-3,
    权重衰减=1e-4,
    训练数据加载器=train_loader,
    验证数据加载器=valid_loader,
    最大轮次=max_epochs,
    gpus=1,
    是否导出ONNX=False,          
    快速开发运行=fast_dev_run   
)

evaluator.拟合(最终模型)

注意

当``fast_dev_run``关闭时，我们在训练100轮后，达到了89.69%的验证准确率。

在DARTS论文中复现结果

通过使用一次探索+再训练策略的简要概述，我们填补了我们的结果（89.69%）与DARTS论文中结果之间的差距。
这是因为我们没有引入一些额外的训练技巧，包括DropPath、
辅助损失、梯度裁剪以及像Cutout那样的数据增强。
他们还训练了更深（20个cell）和更宽（36个filter）的网络，时间更长（600轮）。
在这里，我们重现这些技巧，以获得与DARTS论文相当的结果。

评估器

为了实现这些技巧，我们首先需要对评估器的一些部分进行重写。

使用一次探索策略工作时，评估器需要以PyTorch-Lightning <lightning-evaluator>的风格进行实现，
完整的教程可以在:doc:/nas/evaluator中找到。
简而言之，编写新的评估器的核心部分是编写一个新的LightningModule。
LightingModule_ 是
PyTorch-Lightning中的一个概念，将模型训练过程组织成一系列函数，比如training_step、validation_step、configure_optimizers等。
由于我们仅仅是在:class:~nni.nas.evaluator.pytorch.Classification中添加了一些元素，
我们可以简单地继承:class:~nni.nas.evaluator.pytorch.ClassificationModule，
这是:class:~nni.nas.evaluator.pytorch.Classification背后的基础LightningModule。
一开始可能会感到有些难以理解，但其中大部分只是插拔式的技巧，无需详细了解细节。

import torch
from nni.nas.evaluator.pytorch import ClassificationModule

class DartsClassificationModule(ClassificationModule):
    def __init__(
        self,
        学习率: float = 0.001,
        权重衰减: float = 0.,
        辅助损失权重: float = 0.4,
        最大轮次: int = 600
    ):
        self.辅助损失权重 = 辅助损失权重
        # 训练长度将用于LR调度器
        self.最大轮次 = 最大轮次
        super().__init__(学习率=学习率, 权重衰减=权重衰减, 是否导出ONNX=False)

    def configure_optimizers(self):
        """使用动量自定义优化器，以及一个调度器。"""
        优化器 = torch.optim.SGD(
            self.parameters(),
            动量=0.9,
            lr=self.hparams.学习率,
            weight_decay=self.hparams.权重衰减
        )
        return {
            'optimizer': 优化器.
            'lr_scheduler': torch.optim.lr_scheduler.CosineAnnealingLR(优化器, self.最大轮次, eta_min=1e-3)
        }

    def training_step(self, batch, batch_idx):
        """自定义的训练步骤中带有辅助损失。"""
        x, y = batch
        if self.辅助损失权重:
            y_hat, y_aux = self(x)
            主损失 = self.criterion(y_hat, y)
            辅助损失 = self.criterion(y_aux, y)
            self.log('train_loss_main', 主损失)
            self.log('train_loss_aux', 辅助损失)
            损失 = 主损失 + self.辅助损失权重 * 辅助损失
        else:
            y_hat = self(x)
            损失 = self.criterion(y_hat, y)
        self.log('train_loss', 损失, prog_bar=True)
        for name, metric in self.metrics.items():
            self.log('train_' + name, metric(y_hat, y), prog_bar=True)
        return 损失

    def on_train_epoch_start(self):
        # 在每个轮次之前设置路径丢弃概率。如果模型中未启用路径丢弃，这将不会起作用。
        self.model.set_drop_path_prob(self.model.drop_path_prob * self.current_epoch / self.最大轮次)

        # 在每个轮次开始时记录学习率
        self.log('lr', self.trainer.optimizers[0].param_groups[0]['lr'])

下面是完整的评估器，
它只是将所有内容（当然不包括模型空间和搜索策略）包装在一个对象中。
这里的Lightning是一种特殊类型的评估器。
这里别忘了使用专门用于搜索的训练/验证数据划分（1:1）。

from nni.nas.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
)

策略

创建一个具有梯度裁剪的:class:~nni.nas.strategy.DARTS策略。
如果你熟悉PyTorch-Lightning，你可能意识到梯度裁剪可以在Lightning训练器中启用。
然而，在上面的训练器中启用梯度裁剪不起作用，因为:class:~nni.nas.strategy.DARTS策略的底层实现是基于
手动优化_的。

strategy = Darts策略(梯度剪切值=5.)

启动实验

然后我们使用新创建的评估器和策略再次启动实验。

警告

必须重新实例化``模型空间``，因为有一个已知的限制，即一个模型空间实例不能在多个实验中复用。

model_space = DartsSpace(width=16, num_cells=8, dataset='cifar')

experiment = NasExperiment(model_space, evaluator, strategy)
experiment.run()

exported_arch = experiment.export_top_models(formatter='dict')[0]

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-4.md

exported_arch


当将fast_dev_run设置为False时，我们得到如下架构。在P100 GPU上大约需要8小时。

```python
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
})

重新训练

在重新训练过程中，我们对原始的dataloader进行扩展，引入了另一个被称为Cutout的技巧。Cutout是一种数据增强技术，它会随机遮挡图像中的矩形区域。在CIFAR-10数据集中，通常遮挡的大小为16x16（数据集中的图像大小为32x32）。

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,
])

需要使用新的transform来重新实例化train dataloader，而valid dataloader不受影响，因此可以重复使用。

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)

接下来，根据新的导出架构创建最终模型。这次，启用了辅助损失和drop path概率。

按照论文的相同步骤，我们还将滤波器数量增加到36，将单元数量增加到20，以合理增加模型大小并提升性能。

with model_context(exported_arch):
    final_model = DartsSpace(width=36, num_cells=20, dataset='cifar', auxiliary_loss=True, drop_path_prob=0.2)

创建新的评估器用于重新训练过程，在训练器的关键字参数中加入了梯度裁剪。

max_epochs = 600

evaluator = Lightning(
    DartsClassificationModule(0.025, 3e-4, 0.4, max_epochs),
    trainer=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)

当fast_dev_run关闭时，经过重新训练后，此架构的top-1准确率达到97.12%。如果我们选择整个重新训练过程中的最佳快照，有可能使top-1准确率达到97.28%。

在图中，橙线是训练600个epoch后的验证准确率曲线，而红线对应于在添加所有训练技巧之前的教程中的上一个版本，只训练了100个epoch。

该结果优于DARTS_论文中的"DARTS（一阶）+ cutout"，其准确率仅为97.00±0.14%。尽管我们没有实现二阶版本的DARTS，但其性能与论文中的"DARTS（二阶）+ cutout"是可以相媲美的。我们将在未来计划中实现二阶DARTS，并欢迎您的贡献。

posted @ 2023-12-02 16:59 jasonzhangxianrong 阅读(167) 评论(0) 编辑收藏举报

刷新页面返回顶部

DARTS搜索空间

在DARTS搜索空间中搜索

使用预搜索的DARTS模型

注意

使用DARTS模型空间

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-1.md

注意

在模型空间上搜索

警告

评估器

启动实验

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-2.md

警告

重新训练搜索到的模型

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-3.md

注意

在DARTS论文中复现结果

评估器

策略

启动实验

警告

翻译 private_upload\default\2023-12-02-17-42-43\darts.md.part-4.md

重新训练

公告