NAS

介绍

ModelOpt 提供了一种NAS方法（也称为mode）- 通过modelopt.torch.nas模块的AutoNAS。给定一个模型，这些方法从您提供的基础模型中找到一个满足给定部署约束（例如FLOPs、参数）的子网，几乎不会降低准确性（取决于模型大小减少的激进程度）。关于这种NAS模式的更多详情如下：

autonas: 一种适用于计算机视觉模型的NAS方法，它搜索层级的参数，如通道数量、核大小、网络深度等。

按照以下步骤使用modelopt.torch.nas来获取满足您独特需求的最佳模型：

通过 mtn.convert 转换您的模型：使用一组简单的配置，从您的PyTorch基础模型原生生成神经架构搜索空间。在此过程中，方便地保存和恢复模型架构和权重。
NAS训练：在您现有的训练管道中无缝训练生成的搜索空间。
通过子网架构搜索 mtn.search：搜索满足您部署约束的最佳神经架构（子网），例如， FLOPs / 参数。
微调: 可选地，微调生成的子网以实现更高的准确性。

要了解更多关于NAS及相关概念的信息，请参考以下部分 NAS 概念。

转换并保存

你可以使用mtn.convert()转换你的模型并从中生成一个搜索空间。生成的搜索空间应使用mto.save()保存。可以使用mto.restore()将其加载回来，以执行架构搜索的后续步骤。

示例用法：

import modelopt.torch.nas as mtn
import modelopt.torch.opt as mto
from torchvision.models import resnet50

# User-defined model
model = resnet50()

# Generate the search space for AutoNAS
model = mtn.convert(model, mode="autonas")

# Save the search space for future use
mto.save(model, "modelopt_model.pth")

注意

NAS API 是剪枝 API 的超集。你也可以在这里使用剪枝模式（例如 "fastnas", "gradnas" 等）。

注意

在上面的例子中，我们使用了默认的AutoNAS config 用于 mtn.convert()。你可以使用 mtn.config.AutoNASConfig() 查看它。你也可以指定自定义配置以拥有不同的搜索空间。更多信息请参阅 mtn.convert() 文档。下面显示了一个配置示例：

import modelopt.torch.nas as mtn

config = mtn.config.AutoNASConfig()
config["nn.Conv2d"]["*"]["out_channels_ratio"] += (0.1,)  # include more channel choices

model = mtn.convert(model_or_model_factory, mode=[("autonas", config)])

注意

如果你想了解更多关于转换过程以及模型的先决条件，你可以查看NAS模型先决条件。

注意

请参阅保存和恢复ModelOpt修改的模型以了解所有可用的保存和恢复选项。

分析搜索空间

搜索空间可用于根据您所需的部署约束执行架构搜索。

为了更好地了解性能以及生成的搜索空间的范围，您可以使用 mtn.profile() 来分析搜索空间以及您的部署约束：

import torch

# Looking for a subnet with at most 2 GFLOPs
constraints = {"flops": 2.0e9}

# Measure FLOPs against dummy_input
# Can be provided as a single tensor or tuple of input args to the model.
dummy_input = torch.randn(1, 3, 224, 224)

is_sat, search_space_stats = mtn.profile(model, dummy_input, constraints=constraints)

以下信息将被打印：

Profiling the following subnets from the given model: ('min', 'centroid', 'max').
--------------------------------------------------------------------------------

                            Profiling Results
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Constraint   ┃ min          ┃ centroid     ┃ max          ┃ max/min ratio ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ flops        │ 487.92M      │ 1.84G        │ 4.59G        │ 9.40          │
│ params       │ 4.84M        │ 12.33M       │ 25.50M       │ 5.27          │
└──────────────┴──────────────┴──────────────┴──────────────┴───────────────┘

            Constraints Evaluation
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓
┃              ┃              ┃ Satisfiable  ┃
┃ Constraint   ┃ Upper Bound  ┃ Upper Bound  ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩
│ flops        │ 2.00G        │ True         │
└──────────────┴──────────────┴──────────────┘

Search Space Summary:
----------------------------------------------------------------------------------------------------
* conv1.out_channels                                                               [32, 64]
  conv1.in_channels                                                                [3]
  bn1.num_features                                                                 [32, 64]
* layer1.depth                                                                     [1, 2, 3]
* layer1.0.conv1.out_channels                                                      [32, 64]
  layer1.0.conv1.in_channels                                                       [32, 64]
  layer1.0.bn1.num_features                                                        [32, 64]
* layer1.0.conv2.out_channels                                                      [32, 64]
  ...
  ...
  ...
* layer4.2.conv1.out_channels                                                      [256, 352, 512]
  layer4.2.conv1.in_channels                                                       [2048]
  layer4.2.bn1.num_features                                                        [256, 352, 512]
* layer4.2.conv2.out_channels                                                      [256, 352, 512]
  layer4.2.conv2.in_channels                                                       [256, 352, 512]
  layer4.2.bn2.num_features                                                        [256, 352, 512]
  layer4.2.conv3.out_channels                                                      [2048]
  layer4.2.conv3.in_channels                                                       [256, 352, 512]
----------------------------------------------------------------------------------------------------
Number of configurable hparams: 40
Total size of the search space: 1.90e+18
Note: all constraints can be satisfied within the search space!

你也可以跳过constraints参数，仅打印可用约束的范围，而不检查它是否在你的约束范围内。分析结果将帮助你理解搜索空间，并提出一个可以迭代的潜在搜索约束。

NAS训练

先决条件

在NAS训练期间，您可以使用现有的训练基础设施。然而，我们建议您对训练超参数进行以下修改：

将训练时间（epochs）增加2-3倍。
确保学习率计划根据更长的训练时间进行调整。
我们建议使用连续衰减的学习率计划，例如余弦退火计划（参见PyTorch文档）。

恢复搜索空间

请从保存的搜索空间中恢复，以继续执行其余步骤，如下所示：

# Provide the model before conversion to mto.restore
model = mto.restore(model_or_model_factory, "modelopt_model.pth")

培训

您现在可以继续使用现有的训练流程，同时调整训练时间和学习率。

子网架构搜索

NAS的下一步是在生成的搜索空间上执行架构搜索，以找到满足您部署约束的最佳子网。

先决条件

要在训练好的模型上执行搜索（mtn.search()），需要一个评分函数、一个虚拟输入（用于测量您的部署约束）、训练数据加载器（用于校准归一化层）以及约束条件。请参阅mtn.search() API 以获取更多详细信息。
根据算法的不同，您可能能够通过为每个约束指定上限来提供多个搜索约束，例如flops或params。

执行搜索

以下是在AutoNAS转换和训练模型上运行搜索的示例。

# Wrap your original validation function to only take the model as input.
# This function acts as the score function to rank models.
def score_func(model):
    return validate(model, val_loader, ...)


# Specify the sample input including target data shape for FLOPs calculation.
dummy_input = torch.randn(1, 3, 224, 224)

# Looking for a subnet with at most 2 GFLOPs
search_constraints = {"flops": 2.0e9}

# search_res (dict) contains state_dict / stats of the searcher
searched_model, search_res = mtn.search(
    model=model,
    constraints=search_constraints,
    dummy_input=dummy_input,
    config={
        "data_loader": train_loader,  # training data is used for calibrating BN layers
        "score_func": score_func,  # validation score is used to rank the subnets
        # checkpoint to store the search state and resume or re-run the search with different constraint
        "checkpoint": "modelopt_search_checkpoint.pth",
    },
)

# Save the searched model for further fine-tuning
mto.save(searched_model, "modelopt_searched_model.pth")

提示

如果评分函数的运行时间超过几分钟，考虑对评分函数中使用的数据集进行子采样。可以使用 torch.utils.data.Subset 对PyTorch数据集进行子采样，如下所示：

subset_dataset = torch.utils.data.Subset(dataset, indices)

注意

NAS 将就地修改模型。

注意

mtn.search() 支持通过 PyTorch 中的 DistributedDataParallel 实现分布式数据并行。

微调

搜索后，与剪枝相比，准确率下降可能不那么显著，但我们仍然建议运行微调以恢复最佳准确率。对于AutoNAS，通常一个好的微调计划是重复预训练计划（1x epochs），并使用0.5x-1x的初始学习率，如FastNAS中所做的那样。更多详情请参考剪枝微调部分。

NAS 模型先决条件

在本指南中，我们将逐步介绍如何设置您的模型以与NAS和剪枝一起使用。在本指南结束时，您将能够convert您自己的模型，以生成可用于NAS和剪枝的搜索空间。

转换您的模型

大多数PyTorch模型，包括自定义模型，都与ModelOpt原生兼容（取决于forward的实现方式）。要快速测试您的模型是否兼容，您可以简单地尝试转换它：

import modelopt.torch.nas as mtn
from torchvision.models import resnet50

# User-defined model
model = resnet50()

# Convert the model into a search space
model = mtn.convert(model, mode="fastnas")

如果您遇到问题或想了解更多关于转换过程的信息，请继续阅读。否则，您可以跳过本指南的其余部分。

转换过程

ModelOpt 将自动从您的自定义 PyTorch 模型中为您生成搜索空间。这是在剪枝和神经架构搜索（NAS）期间执行的一次性过程。一旦模型被转换，您可以保存并恢复它，用于下游任务，如训练、推理和微调。

为了帮助您更好地理解搜索空间是如何从您的模型中得出的，我们将在下面更详细地介绍这个过程。

层支持

您可以通过使用由层组成的模型架构来充分利用ModelOpt，这些层可以被ModelOpt自动转换为可搜索单元。

具体来说，以下原生PyTorch层可以转换为可搜索单元：

import torch.nn as nn

# We convert native PyTorch convolutional layers to automatically search over the number of
# channels and optionally over the kernel size.
nn.Conv1d
nn.Conv2d
nn.Conv3d
nn.ConvTranspose1d
nn.ConvTranspose2d
nn.ConvTranspose3d

# We convert native PyTorch linear layers to automatically search over the number of features
nn.Linear

# We convert native PyTorch sequential layers that contain residual blocks to automatically
# search over the number of layers (depth) in the sequential layer.
nn.Sequential

# We convert Megatron-core / NeMo GPT-style models (e.g. Llama3.1, NeMo Mistral, etc.)
# to automatically search over the MLP hidden size, number of attention heads, number of GQA groups,
# and depth of the model.
megatron.core.transformer.module.MegatronModule
nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTModel

# We convert Hugging Face Attention layers to automatically search over the number of heads
# and MLP hidden size.
# Make sure `config.use_cache` is set to False during pruning.
transformers.models.bert.modeling_bert.BertAttention
transformers.models.gptj.modeling_gptj.GPTJAttention

生成搜索空间

要从您想要的模型生成搜索空间，只需调用 mtn.convert() 即可：

import modelopt.torch.nas as mtn
from torchvision.models import resnet50

# User-defined model
model = resnet50()

# Convert the model for NAS/pruning
model = mtn.convert(model, mode="fastnas")

您生成的model代表了一个由多个子网组成的搜索空间。请注意，您可以像使用其他常规的PyTorch模型一样使用转换后的模型。它将根据当前激活的子网进行行为。

大致上，convert 过程可以分为以下步骤：

通过模型进行追踪，以解决层依赖关系并记录层之间的连接方式。
将支持的层转换为可搜索的单元，即动态层，并根据记录的依赖关系连接它们。
从转换后的模型生成一致的搜索空间。

注意

在剪枝过程中，当调用mtp.prune时，转换会隐式执行。

先决条件

为了正确生成搜索空间，您的原始模型应满足以下先决条件。

可追溯性

模型需要能够使用ModelOpt的torch.fx类似的追踪器进行追踪。

如果没有，当你运行mtn.convert()时，你会看到错误或警告。请注意，其中一些警告可能不会影响搜索空间，因此可以忽略。

请注意，在某些情况下，某些层无法被追踪，如果可能的话，您应该调整它们的定义和前向方法以使其可追踪。否则，在转换过程中，这些层以及所有受影响的层将被忽略。

分布式数据并行

在转换过程之后，应该用DistributedDataParallel包装模型，并且在包装过程中需要设置find_unused_parameters=True：

model = mtn.convert(model, ...)
model = DistributedDataParallel(model, find_unused_parameters=True)

辅助模块

如果你的模型包含辅助模块，例如仅在训练期间活动的分支，请确保转换整个模型，以便在转换过程中所有模块都处于活动状态。

已知限制

请注意NAS 常见问题解答中提到的其他潜在限制！

NAS 概念

下面，我们将概述ModelOpt的神经架构搜索（NAS）和剪枝算法以及其基本概念和术语。

概述

Glossary
神经架构搜索 (NAS)	为给定任务寻找最佳神经网络架构的过程。
搜索空间	在剪枝或NAS期间搜索的可能候选架构集合。
架构超参数	描述搜索空间的一组超参数，例如层数。
子网	搜索空间中的一个候选架构。
基于NAS的训练	在搜索空间中训练子网集合的过程。
架构搜索与选择	在训练好的搜索空间中找到最优子网的过程。
子网微调	为了提高最终准确性，单独训练所选子网的过程。
Pruning	从神经网络中移除给定任务的冗余组件的过程。

概念

下面，我们介绍神经架构搜索的概念和术语。在常规的神经网络训练中，只有神经网络的权重被训练。然而，在NAS中，模型的权重和架构都被训练。

神经架构搜索 (NAS)

神经架构搜索是从一组候选架构中找到最佳神经网络架构的过程。NAS通常在训练之前、期间或之间进行。在NAS过程中，使用不同的性能指标（如准确性、设备上的延迟或模型的大小）来评估候选架构。

搜索空间

搜索空间被定义为所有可能训练的神经架构的（离散）集合。搜索空间来源于（用户指定的）基础架构（例如，ResNet50）和一组描述如何参数化基础架构的配置，更多信息请参见NAS模型先决条件。

架构超参数

搜索空间通过一组离散的架构超参数进行参数化，这些超参数描述了基础架构的各个“修改”，例如卷积层中的通道数量、重复构建块的数量、变压器层中的注意力头数量等。搜索空间中的每个可能架构都可以描述为架构超参数集的不同配置。

子网

搜索空间由一组子网组成，其中每个子网代表一个神经架构。每个子网构成一个具有不同层和操作符或每层不同参数化（例如通道数）的神经架构。

为了更好地描述给定的搜索空间，我们通常考虑几个不同的子网：

最小子网 (min): 搜索空间内的最小子网。
中心子网 (centroid): 每个架构超参数设置为最接近其中心（平均值）的子网。
最大子网 (max): 搜索空间内最大的子网。

ModelOpt转换后的模型

转换后，用户提供的神经网络将代表搜索空间。可以通过mtn.convert()获取，参见转换和保存。

在转换过程中，搜索空间会自动从给定的基础架构中派生，并且相关的架构超参数会自动识别。

下一步是训练转换后的模型（而不是原始架构），以找到适合您部署约束的最佳子网。

基于NAS的训练

在搜索空间的训练过程中，我们同时训练模型的权重和架构：

使用 modelopt.torch.nas 你可以重用你现有的训练循环来训练搜索空间。
在搜索空间训练期间，整个子网集合与其权重一起自动训练。
鉴于我们同时训练架构（所有子网络）和权重，训练数据可能与上述NAS训练部分中描述的常规训练有所不同。

架构搜索与选择

在搜索空间训练过程结束时，下一步是从搜索空间中搜索并选择子网：

搜索过程是一个离散优化问题，用于从搜索空间中确定最佳子网配置。
搜索过程会考虑您的部署约束，例如FLOPs、参数或延迟和推理设备，以确定在满足约束条件的同时最优（最准确）的子网配置。
生成的子网可以用于进一步的下游任务，例如微调和部署。

子网微调

为了进一步提高所选子网的准确性，通常会在原始任务上对子网进行微调：

要微调子网，您可以简单地重复原始模型的训练流程，使用上述微调部分中描述的调整后的训练计划。
微调后的模型构成了可部署模型，它在准确性和您提供的约束之间实现了最佳权衡。

NAS 与剪枝

NAS和剪枝的区别总结如下。

	NAS	剪枝
搜索空间	更灵活的搜索空间，增加了可搜索的维度，如网络深度、核大小或激活函数的选择。	搜索空间灵活性较低，可搜索的维度被限制在较少的选项上，例如通道数量、特征数量或注意力头数量。
训练时间	通常需要在搜索子网之前额外训练模型一段时间。	当有预训练的检查点可用时，不需要进行训练。如果没有，可以使用常规训练来预训练一个检查点。
性能	由于更灵活的搜索空间和增加的训练时间，可以提供更好的准确性与延迟的权衡。	在某些特定应用中，可能提供与NAS相似的性能，然而，由于搜索空间和训练时间的限制，通常表现出较差的性能。

NAS

介绍

转换并保存

分析搜索空间

NAS训练

先决条件

恢复搜索空间

培训

子网架构搜索

先决条件

执行搜索

微调

NAS 模型先决条件

转换您的模型

转换过程

层支持

生成搜索空间

先决条件

可追溯性

分布式数据并行

辅助模块

已知限制

NAS 概念

概述

概念

神经架构搜索 (NAS)

搜索空间

架构超参数

子网

ModelOpt转换后的模型

基于NAS的训练

架构搜索与选择

子网微调

NAS 与 剪枝

NAS 与剪枝