关于AI 神经网络架构搜索的几点记录

神经架构搜索 NAS 架构搜索 自动化 深度学习
作者: 醉月思 发布日期: 2023-08-11 分类: 技术实践

项目背景其实挺简单：给一个医疗影像分类任务做模型优化，数据量约 10 万张，算力限制在单张 A100 GPU，目标是在推理时间不超过 50ms 的前提下把准确率从 85% 提升到 88%。

这种做法在过去几年里确实管用，但这次卡在 86% 左右就死活上不去了。

起因：为什么需要 NAS

项目背景其实挺简单：给一个医疗影像分类任务做模型优化，数据量约 10 万张，算力限制在单张 A100 GPU，目标是在推理时间不超过 50ms 的前提下把准确率从 85% 提升到 88%。

按照传统思路，我会先在几个主流架构里选一个，比如 ResNet、EfficientNet、MobileNetV3，然后基于这个架构做微调：调整深度、宽度、分辨率、激活函数、正则化强度。这种做法在过去几年里确实管用，但这次卡在 86% 左右就死活上不去了。

问题在于，手动调参有几个绕不开的限制：

  1. 搜索空间有限：我只会在自己熟悉的架构里选，比如 ResNet 的三种残差连接类型、EfficientNet 的复合缩放系数。这个搜索空间实际上很小，容易陷入局部最优。

  2. 经验主义陷阱：我相信残差连接能缓解梯度消失，所以倾向于在每三层之间加一个；相信深度可分离卷积能减少参数量，所以把标准卷积都替换掉。这些信念未必在当前任务上成立。

  3. 时间成本高：每次调整架构都要重新训练好几天，试错成本太高。理论上应该做更系统的网格搜索或贝叶斯优化，但实际上根本等不起。

  4. 硬件约束难量化：推理时间不超过 50ms 这个限制，手动设计架构时只能靠感觉，比如"这个模型参数量少应该就快"，但实际上显存带宽、计算密集度、并行度都影响最终速度。

这些问题不是第一次遇到，但这次实在不想再靠经验瞎猜了，于是决定试试 NAS。

NAS 的基本思路

简单来说，NAS 就是把神经网络架构设计问题当成搜索问题来解。搜索空间里包含各种可能的架构设计，比如每一层用什么类型的卷积、几个通道、激活函数选什么；搜索策略在这个空间里找最优解；评估策略判断一个架构好不好听。

整个过程可以抽象成三个部分：

graph LR A[搜索空间定义] --> B[搜索策略] B --> C[评估策略] C --> B B --> D[最优架构]

搜索空间定义要回答"我们允许哪些操作"。比如在卷积网络里，每一层可能是 3x3 卷积、5x5 卷积、深度可分离卷积、空洞卷积、池化层、跳跃连接，也可能是 Identity。每一层可以选的操作越多，搜索空间越大，找到更好架构的理论上可能性越高，但搜索成本也成指数增长。

搜索策略要解决"怎么在搜索空间里找"。最笨的办法是随机搜索，但效率太低；稍聪明一点用进化算法，种群里的个体是不同的架构，按精度和延迟做选择、交叉、变异；更复杂一点用强化学习，让一个 RNN 学习生成架构，用准确率做奖励。

评估策略是另一个瓶颈。最准确的方法是每个候选架构都从头训练完再评估精度，但这样太慢。后来有了权重共享技术，让所有候选架构共享一部分预训练权重，不需要每个架构都从头训练；还有基于可微分架构搜索的方案，把离散选择问题改成连续优化问题。

选择什么 NAS 方案

理论上 NAS 方案很多，但实际项目中要考虑几个现实约束：算力有限、时间有限、不能保证一定能超过手动设计的架构、工具链要足够成熟。

在查阅了大量论文和开源项目后，我选择了 EfficientNet + DARTS 的组合路线：

  • EfficientNet 作为基线，因为它的复合缩放思想在资源受限场景下表现不错
  • DARTS（Differentiable Architecture Search）作为搜索策略，因为它相对于强化学习和进化算法效率高很多，且实现相对成熟
  • 权重共享加上可微分搜索，把搜索过程压缩到几小时级别

这个选择不是没有争议。DARTS 在某些任务上被指出存在稳定性问题，搜索出来的架构在某些情况下不如基线。但考虑到我们的算力和时间限制，这是性价比最高的选择。

实现细节

搜索空间定义

先定义一个基础的搜索空间，每个候选操作包括以下几种：

# ops.py
import torch
import torch.nn as nn

OPS = {
    'none': lambda C, stride, affine: Zero(stride),
    'avg_pool_3x3': lambda C, stride, affine: nn.AvgPool2d(3, stride=stride, padding=1, count_include_pad=False),
    'max_pool_3x3': lambda C, stride, affine: nn.MaxPool2d(3, stride=stride, padding=1),
    'skip_connect': lambda C, stride, affine: Identity() if stride == 1 else FactorizedReduce(C, C, affine=affine),
    'sep_conv_3x3': lambda C, stride, affine: SepConv(C, C, 3, stride, 1, affine=affine),
    'sep_conv_5x5': lambda C, stride, affine: SepConv(C, C, 5, stride, 2, affine=affine),
    'dil_conv_3x3': lambda C, stride, affine: DilConv(C, C, 3, stride, 2, 2, affine=affine),
    'dil_conv_5x5': lambda C, stride, affine: DilConv(C, C, 5, stride, 4, 2, affine=affine),
}

class SepConv(nn.Module):
    def __init__(self, C_in, C_out, kernel_size, stride, padding, affine=True):
        super(SepConv, self).__init__()
        self.op = nn.Sequential(
            nn.ReLU(inplace=False),
            nn.Conv2d(C_in, C_in, kernel_size, stride, padding, groups=C_in, bias=False),
            nn.Conv2d(C_in, C_out, 1, 1, 0, bias=False),
            nn.BatchNorm2d(C_out, affine=affine),
            nn.ReLU(inplace=False),
            nn.Conv2d(C_out, C_out, kernel_size, 1, padding, groups=C_out, bias=False),
            nn.Conv2d(C_out, C_out, 1, 1, 0, bias=False),
            nn.BatchNorm2d(C_out, affine=affine),
        )

    def forward(self, x):
        return self.op(x)

这个搜索空间包含了 8 种操作，覆盖了常见的卷积和池化类型。理论上可以再加更多，但这样搜索空间会指数级增长，我们的算力扛不住。

网络结构

定义一个包含 normal cell 和 reduction cell 的网络结构：

# model.py
import torch.nn as nn
import torch.nn.functional as F

class Network(nn.Module):
    def __init__(self, C, num_classes, layers, criterion, steps=4, multiplier=4, stem_multiplier=3):
        super(Network, self).__init__()
        self._C = C
        self._num_classes = num_classes
        self._layers = layers
        self._criterion = criterion
        self._steps = steps
        self._multiplier = multiplier

        C_curr = stem_multiplier * C
        self.stem = nn.Sequential(
            nn.Conv2d(3, C_curr, 3, padding=1, bias=False),
            nn.BatchNorm2d(C_curr)
        )

        C_prev_prev, C_prev, C_curr = C_curr, C_curr, C
        self.cells = nn.ModuleList()
        reduction_prev = False
        for i in range(layers):
            if i in [layers // 3, 2 * layers // 3]:
                C_curr *= 2
                reduction = True
            else:
                reduction = False
            cell = Cell(steps, multiplier, C_prev_prev, C_prev, C_curr, reduction, reduction_prev)
            reduction_prev = reduction
            self.cells += [cell]
            C_prev_prev, C_prev = C_prev, multiplier * C_curr

        self.global_pooling = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(C_prev, num_classes)

    def forward(self, input, target):
        s0 = s1 = self.stem(input)
        for i, cell in enumerate(self.cells):
            s0, s1 = s1, cell(s0, s1, self.drop_path_prob)
        out = self.global_pooling(s1)
        logits = self.classifier(out.view(out.size(0), -1))
        return self._criterion(logits, target) * 0.0  # 为了兼容，实际损失在训练循环中计算

    def loss(self, input, target):
        logits = self(input, target)
        return self._criterion(logits, target)

这个结构参考了 DARTS 的原始实现，用 normal cell 做特征提取，用 reduction cell 做下采样。

搜索过程

DARTS 的核心思想是把架构搜索问题变成连续优化问题。每个候选操作都有一个对应的权重参数，训练时通过 softmax 把权重归一化，得到每个操作被选中的概率：

# search_train.py
import torch
import torch.optim as optim
from model import Network
from ops import OPS

def train(search_loader, valid_loader, model, architect, optimizer, args):
    for epoch in range(args.epochs):
        lr = scheduler.get_lr()[0]
        model.drop_path_prob = args.drop_path_prob * epoch / args.epochs

        # 训练网络权重
        train_acc, train_obj = train_epoch(search_loader, model, architect, optimizer, lr)

        # 更新架构参数
        valid_acc, valid_obj = infer(valid_loader, model)

        print(f'Epoch {epoch}: train_acc {train_acc:.4f} train_obj {train_obj:.4f} valid_acc {valid_acc:.4f} valid_obj {valid_obj:.4f}')

    # 提取最终架构
    genotype = model.genotype()
    print(f'Final genotype: {genotype}')

def train_epoch(train_loader, model, architect, optimizer, lr):
    model.train()
    architect.optimizer.zero_grad()

    for step, (input, target) in enumerate(train_loader):
        input = input.cuda()
        target = target.cuda(non_blocking=True)

        # 架构参数的梯度计算
        input_search, target_search = next(iter(valid_loader))
        input_search = input_search.cuda(non_blocking=True)
        target_search = target_search.cuda(non_blocking=True)

        architect.step(input, target, input_search, target_search, lr, optimizer, unrolled=args.unrolled)

        # 网络权重的训练
        optimizer.zero_grad()
        logits = model(input)
        loss = criterion(logits, target)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        optimizer.step()

这里的关键是分离了网络权重和架构参数的优化。网络权重用训练集优化，架构参数用验证集优化，这样避免过拟合到训练数据。

从搜索空间到具体架构

搜索完成后，从每个节点的候选操作中选择权重最大的那个，构建最终的网络：

# extract_arch.py
def extract_final_architecture(model, genotype):
    final_architecture = []

    # 遍历每个 cell
    for cell_type in ['normal', 'reduce']:
        cell_ops = genotype[cell_type]

        for op_idx, (op_name, prev_nodes) in enumerate(cell_ops):
            final_architecture.append({
                'cell_type': cell_type,
                'op_index': op_idx,
                'operation': op_name,
                'input_nodes': prev_nodes
            })

    return final_architecture

这个提取出来的架构就是 NAS 搜索的结果，可以用于后续的训练和部署。

踩过的坑

坑一：搜索不稳定

刚开始搜索的时候，发现不同次运行的结果差异很大。有时候能搜到比 ResNet 好的架构，有时候甚至不如基线。

查了一圈发现几个问题：

  1. 学习率太激进：原始 DARTS 论文建议架构参数和网络权重都用同一个学习率，但实际操作中发现架构参数的学习率应该更小一些，否则容易震荡。

  2. 权重初始化敏感：不同次运行的随机种子不同，导致初始化的权重分布不同，进而影响搜索结果。后来固定了随机种子，并用了更保守的初始化策略。

  3. 搜索空间过大：原本想尝试更复杂的搜索空间，比如加上 attention、group convolution 等操作，结果搜索过程直接崩了。最后还是缩减到基础的卷积和池化操作。

修复这些问题后，搜索结果的稳定性明显提升，多次运行的结果差异在可接受范围内。

坑二：搜索出来的架构训练不收敛

搜到一个看起来很不错的架构，但从头训练时一直不收敛，准确率在 50% 左右晃悠。

排查后发现几个原因：

  1. 深度太深：搜索过程中共享了权重，所以训练过程相当于在部分预训练的模型上做微调。但最终架构需要从头训练，太深的网络容易出现梯度消失。

  2. 跳跃连接过多：搜索出来的架构倾向于大量使用 skip_connect，这在搜索过程中有优势，但训练时容易出现梯度爆炸或消失。

  3. 激活函数选择不当：某些架构在不同层之间混用了不同激活函数，导致梯度分布不均匀。

解决方案是加了几项约束：

# 添加架构约束
def apply_architecture_constraints(genotype):
    # 限制最大深度
    if count_max_depth(genotype) > 50:
        return None

    # 限制跳跃连接比例
    skip_count = count_skip_connections(genotype)
    total_count = count_total_connections(genotype)
    if skip_count / total_count > 0.4:
        return None

    # 统一激活函数
    if has_mixed_activation(genotype):
        return None

    return genotype

加上这些约束后，搜索出来的架构更容易训练收敛。

坑三：推理时间不可控

最初搜索时只考虑了精度，没管推理时间。结果搜出来的架构虽然精度不错，但在 A100 上跑完要 80ms，远超 50ms 的限制。

这个问题有几个解决思路：

  1. 在奖励函数中加入延迟惩罚：用 FLOPs 或实际测量的推理时间作为延迟指标，和精度一起组成奖励。

  2. 限制模型大小：在搜索空间里限制每一层的最大通道数，或限制总参数量。

  3. 后期剪枝：先用 NAS 搜索出高精度架构，再用结构化剪枝或量化减少延迟。

最后采用了第一种方案，奖励函数改成：

def compute_reward(accuracy, latency, target_latency=50):
    # 准确率权重
    acc_weight = 1.0

    # 延迟惩罚：超过目标延迟线性惩罚
    if latency <= target_latency:
        latency_penalty = 0.0
    else:
        latency_penalty = (latency - target_latency) / target_latency

    # 最终奖励
    reward = acc_weight * accuracy - 2.0 * latency_penalty
    return reward

延迟惩罚的权重调了多次，2.0 这个值是在多次试验后找出来的平衡点。

坑四：硬件不匹配

搜索过程是在 A100 上完成的，但最终部署的目标设备是 NVIDIA T4，两者的架构差异导致性能表现不一致。

A100 的 Tensor Core 很强，所以搜索出来的架构倾向于使用大量矩阵运算；但 T4 的计算能力弱一些，反而更依赖显存带宽。结果是在 A100 上跑得快的架构，在 T4 上未必最优。

这个问题的解决方案是在搜索过程中加入硬件感知：

# 硬件感知的奖励函数
def compute_hardware_aware_reward(accuracy, latency_a100, latency_t4, target_latency=50):
    # 在两个设备上的延迟都要考虑
    latency_penalty_a100 = max(0, latency_a100 - target_latency) / target_latency
    latency_penalty_t4 = max(0, latency_t4 - target_latency) / target_latency

    # T4 的延迟权重更大，因为它是目标设备
    reward = accuracy - 1.0 * latency_penalty_a100 - 3.0 * latency_penalty_t4
    return reward

但这样会成倍增加搜索成本，因为需要在两个设备上测量延迟。最终折中的方案是只在搜索初期用 A100，后期用 T4 做二次筛选。

最终结果

经过这次折腾，最终得到的架构在 T4 上的推理时间是 48ms，准确率 87.8%，相比最初的 ResNet-50（精度 85.4%，延迟 62ms）有明显提升。

为了对比，还跑了几种手动设计的架构：

架构准确率推理时间参数量
ResNet-5085.4%62ms25.6M
EfficientNet-B386.2%58ms12.2M
MobileNetV3-Large84.9%45ms5.4M
NAS 搜索架构87.8%48ms8.7M

用散点图把各架构的准确率与推理延迟放在同一坐标系，能一眼看出 NAS 结果是否同时满足 50ms 约束和精度目标。

ResNet、EfficientNet、MobileNet 与 NAS 架构在准确率-延迟平面上的分布对比

NAS 搜索出来的架构在精度和延迟之间找到了更好的平衡点，既压进了 50ms 预算，又明显高于手动调参的上限。

写在后面

这次 NAS 实践花了一个多月，从最初的盲目尝试到后来慢慢摸索出适合自己的方案，中间踩过不少坑。

解决了：

  • 经验主义陷阱：不再依赖直觉调参，用搜索策略探索更大的架构空间
  • 硬件约束：在奖励函数中显式考虑延迟，避免搜出无法部署的架构
  • 搜索效率：用权重共享和可微分搜索把搜索过程压缩到可接受范围

留下了：

  • 搜索不稳定：不同次运行的结果仍有差异，需要进一步改进
  • 硬件迁移成本：在不同硬件间迁移时需要重新搜索或调整
  • 工具链复杂度：整个流程涉及多个组件，维护成本高

NAS 不是银弹，但在资源受限的场景下确实能带来价值。前提是你有足够的算力和时间投入，且对搜索结果有合理的预期。

后续还值得尝试的方向：混合精度搜索、神经网络编译器协同优化、跨硬件的架构迁移。但这些得看业务优先级和团队能力。

最后提醒一句：NAS 工具在不断演进，今天看起来复杂的方案，明年可能就有更成熟的框架可以直接用。保持关注，但不要盲目追新。

这次 NAS 实践没有用现成的 AutoML 平台，而是自己实现了基础的 DARTS 流程。虽然花了不少时间，但让我对架构搜索的细节有了更深的理解，也积累了一些可以复用的组件。

版权声明: 本文首发于 指尖魔法屋-关于AI 神经网络架构搜索的几点记录（https://blog.thinkmoon.cn/post/333-ai-neural-architecture-search-manual-search-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章