面向密集预测任务的神经网络架构搜索：从原理到工程实践

1. 项目概述与核心价值

“神经网络架构搜索在密集预测任务中的应用与优化”，这个标题听起来很学术，但背后其实是我们这些在一线搞计算机视觉、图像分割、深度估计的工程师和研究员们每天都在琢磨的“硬骨头”。简单来说，它探讨的是如何让机器自己找到最适合做“像素级”预测任务的网络结构。什么是密集预测？就是给输入图像的每一个像素都打上一个标签或预测一个值，比如语义分割（识别图中每个像素属于“人”、“车”还是“天空”）、深度估计（预测每个像素距离摄像机的远近）、图像增强（对每个像素进行去噪或超分辨率处理）等等。这些任务对模型的精度和效率要求极高，传统上我们依赖专家经验手动设计网络（如U-Net、DeepLab系列），但这个过程耗时耗力，且容易陷入局部最优。

NAS（Neural Architecture Search）的出现，就是为了解决这个痛点：用算法自动搜索出高性能的网络架构。然而，直接把为图像分类设计的NAS方法套用到密集预测任务上，往往会“水土不服”。分类任务关心的是整张图的“是什么”，而密集预测关心的是每个像素的“是什么”或“是多少”，这导致了在搜索空间设计、评估指标、计算开销上存在根本性差异。这个项目要做的，就是深入这些差异，并针对性地进行优化。如果你正在为你的分割模型精度卡在某个瓶颈而烦恼，或者苦恼于模型太大无法部署到边缘设备，那么理解并应用NAS于密集预测，可能就是你的破局点。接下来，我将结合多年的实战经验，拆解这里面的门道、实操步骤以及那些容易踩坑的细节。

2. 核心思路与方案选型背后的考量

为什么不能直接用现成的NAS？这得从密集预测任务的特殊性说起。第一，感受野与多尺度信息。密集预测需要模型同时理解局部细节（如物体边缘）和全局上下文（如物体类别）。这就要求搜索出的架构必须具备高效融合多尺度特征的能力，比如包含丰富的跳跃连接、空洞卷积或特征金字塔模块。而分类NAS搜索出的结构可能更侧重于逐步下采样、压缩信息，这对像素级定位不利。

第二，计算开销的鸿沟。评估一个分类模型，输入一张图，输出一个类别向量，计算量相对固定。但评估一个分割模型，输入一张图，输出一张与原图尺寸相当的分割图，计算量尤其是内存占用（由于要保存高分辨率特征图）会呈数量级增长。许多经典的NAS方法（如基于强化学习或进化算法的）需要评估成千上万个候选架构，这个成本在密集预测任务上是不可承受的。

第三，优化目标的不同。分类常用Top-1准确率，而密集预测任务使用mIoU（平均交并比）、Pixel Accuracy、RMSE（均方根误差）等。这些指标与网络结构的关系更为复杂，且训练过程更不稳定，使得搜索算法的信号更嘈杂。

因此，我们的方案选型必须围绕以下几点展开：

1. 采用权重共享的One-Shot NAS方法：这是当前的主流和务实选择。代表方法有DARTS、ENAS、ProxylessNAS等。其核心思想是构建一个包含所有可能操作的超网络（Supernet），在一次训练中学习所有子架构的权重。搜索时，只需要通过不同的路径采样（即选择不同的操作和连接），就能快速评估子网络性能，避免了重复训练，极大降低了计算成本。对于计算密集型的密集预测任务，这是几乎唯一可行的路径。
2. 设计面向密集预测的搜索空间：这是优化的核心。我们不能简单使用为ImageNet分类设计的搜索空间（如堆叠相同的Cell）。需要将那些在密集预测中被验证有效的模块纳入可搜索选项。例如：
   • 可搜索操作：除了标准的卷积、池化，必须加入空洞卷积（Dilated Conv）以扩大感受野而不损失分辨率；加入可变形卷积（Deformable Conv）以适应不规则形状；考虑注意力机制（如SE、CBAM）的嵌入。
   • 可搜索连接：支持密集跳跃连接（类似U-Net的编码器-解码器跨层连接）、特征金字塔连接（FPN）等。搜索空间需要定义在哪个层级之间建立连接，以及连接的方式（相加、拼接、注意力加权）。
   • 多尺度特征融合单元：设计一个可搜索的“融合单元”，让算法决定如何融合来自骨干网络不同层级的特征图，是上采样后拼接，还是先卷积再相加。
3. 定制化的性能评估策略：由于直接在整个大型数据集（如Cityscapes， 2048x1024分辨率）上评估子网络依然很慢，我们需要策略：
   • 代理任务：在搜索阶段，使用分辨率更低的图片（如256x256）或数据集的一个小子集进行快速评估。
   • 性能预测器：训练一个回归模型，根据子架构的编码（如一个表示架构的向量）来预测其在大任务上的性能，从而替代部分实际评估。
   • 早停机制：对采样的子网络进行少量迭代的训练，根据其早期表现排序，只对排名靠前的进行完整训练和评估。

选择这样的方案，是在搜索质量、计算效率和实现复杂性之间取得的平衡。完全从零开始的搜索（如强化学习）质量可能更高，但成本我们负担不起；而过于简化的搜索空间又无法发挥NAS的潜力。权重共享的One-Shot方法配合精心设计的搜索空间，是目前最具实操性的路线。

3. 面向密集预测的NAS搜索空间设计详解

搜索空间的设计是NAS成功的基石，对于密集预测任务更是如此。一个糟糕的搜索空间，再好的搜索算法也找不出金子。这里，我分享一个我们实践中验证过的、层次化的搜索空间设计思路。

3.1 骨干网络（Backbone）搜索空间

骨干网络负责提取多层次的特征。我们通常将其设计为多个阶段（Stage），每个阶段包含若干可搜索的单元（Cell）。

• 单元（Cell）结构：借鉴DARTS，每个Cell是一个有向无环图，包含N个节点。每个节点代表一个特征图，每条边代表一个可搜索的操作（如3x3 Conv, 5x5 Conv, 3x3 Dilated Conv (rate=2), 5x5 Dilated Conv (rate=2), Identity, Zero）。在密集预测中，我们特别强调加入空洞卷积和可变形卷积作为候选操作。
• 阶段（Stage）间设计：每个Stage之间通过步长为2的卷积或池化进行下采样。我们可以让算法搜索下采样的位置和方式（例如，是使用普通卷积下采样还是使用空洞卷积保持分辨率更久？），但这会大幅增加搜索空间。一个更实用的方法是固定下采样策略，而专注于Cell内部和Cell之间连接方式的搜索。
• 通道数搜索：每一层的输出通道数也是一个重要的超参数。我们可以为每个Stage预设几个宽度乘数选项（如0.5x, 0.75x, 1.0x, 1.25x），让搜索算法选择，以实现精度和速度的权衡。

注意：盲目扩大搜索空间（如加入太多操作或连接选项）会导致超网络训练极其困难，且容易过拟合到搜索验证集上。一个原则是：纳入你认为最有可能有效的几个核心操作，而不是所有已知操作。例如，在初期，可以只包含标准卷积、空洞卷积和Identity。

3.2 特征融合网络（Neck）搜索空间

这是密集预测NAS区别于分类NAS的关键。骨干网络提取了多尺度特征{C1, C2, C3, C4, C5}（分辨率递减），我们需要将它们高效地融合起来，用于最终的像素预测。

我们可以设计一个可搜索的特征金字塔网络。例如，定义一个从深层到浅层的融合过程：

1. P5 = 对C5进行1x1卷积。
2. 对于P4的生成，我们需要融合C4和上采样后的P5。这里就可以引入搜索：
   • 融合方式搜索：候选操作包括Sum（逐元素相加）、Concat（通道拼接后接1x1卷积）、Attention_Fusion（通过注意力权重加权融合）。
   • 上采样方法搜索：候选操作包括Nearest_Upsample（最近邻上采样）、Bilinear_Upsample（双线性上采样）、Transpose_Conv（转置卷积）。
3. 融合后的P4再经过一个可搜索的** refinement 模块**（一个微型的搜索Cell，用于进一步优化特征），然后继续向P3、P2融合。

通过这种方式，搜索算法可以自动发现最适合当前任务的多尺度特征融合路径。例如，对于需要精细边界的医学图像分割，算法可能更倾向于选择能保留更多高频信息的融合方式。

3.3 解码器与预测头搜索空间

解码器负责将融合后的特征图上采样回原图分辨率。这里同样可以引入搜索：

• 上采样路径搜索：是逐步上采样（P2 -> 输出），还是跳跃连接到更底层的特征（如直接融合C1）后再上采样？可以定义几条候选路径让算法选择。
• 预测头结构搜索：最终的预测头可以是一个简单的1x1卷积，也可以是一个小型的多层感知机（MLP）。我们可以让算法搜索预测头的层数和通道数。

3.4 搜索空间的代码化表示

在实际实现中，我们需要用代码将上述搜索空间定义出来。通常使用一个全局的“架构参数”alpha来表示选择概率。以下是一个简化的PyTorch风格示例，展示如何定义融合方式的搜索边：

import torch import torch.nn as nn import torch.nn.functional as F class SearchableFusionLayer(nn.Module): def __init__(self, in_channels_high, in_channels_low, out_channels): super().__init__() # 定义候选操作集合 self.candidate_ops = nn.ModuleDict({ 'sum': nn.Identity(), # 相加操作不需要参数 'concat': nn.Sequential( nn.Conv2d(in_channels_high + in_channels_low, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels) ), 'attention': AttentionFusion(in_channels_high, in_channels_low, out_channels) # 自定义的注意力融合模块 }) # 架构参数，长度等于候选操作数量 self.alpha = nn.Parameter(torch.randn(len(self.candidate_ops))) def forward(self, high_res_feat, low_res_feat): # 上采样高分辨率特征到低分辨率特征的大小 high_res_feat_up = F.interpolate(high_res_feat, size=low_res_feat.shape[2:], mode='bilinear', align_corners=False) # 根据架构参数alpha，计算每个操作的权重（使用Gumbel-Softmax得到可导的离散分布） weights = F.gumbel_softmax(self.alpha, tau=1, hard=False, dim=-1) output = 0 for i, (op_name, op) in enumerate(self.candidate_ops.items()): if op_name == 'sum': feat = op(high_res_feat_up) + low_res_feat # op是Identity elif op_name == 'concat': feat = op(torch.cat([high_res_feat_up, low_res_feat], dim=1)) else: # 'attention' feat = op(high_res_feat_up, low_res_feat) # 加权求和 output += weights[i] * feat return output

在搜索阶段，我们使用weights对各个操作的结果进行加权求和（可导）。在最终架构确定时，我们选择weights中最大的那个索引对应的操作，得到一个确定的、轻量的网络。

4. 基于权重共享超网络的训练与搜索实战

设计好搜索空间后，下一步就是训练超网络并执行搜索。这个过程分为两个阶段：超网络训练和架构搜索。

4.1 超网络训练阶段

这个阶段的目标是训练一个权重共享的超网络，使其内部所有子网络（即不同的架构选择）的权重都得到相对良好的学习。这是整个流程中最耗资源但也最关键的步骤。

训练流程：

1. 前向传播：在每个训练迭代（iteration）中，随机采样一个子架构。采样方法通常是根据当前的架构参数alpha（经过softmax或gumbel-softmax）进行随机采样，或者直接采样概率最大的操作（在训练后期为了稳定）。
2. 损失计算：用采样子架构在训练数据上进行前向传播，计算损失（对于分割任务，常用交叉熵损失和Dice损失的组合）。
3. 反向传播：损失同时对网络权重（W）和架构参数（alpha）求导，并进行更新。
   • 网络权重W：使用标准的优化器（如AdamW）更新。
   • 架构参数alpha：通常使用另一个优化器（如Adam）更新，并且其学习率一般设置得比W的学习率更大（例如10倍），以便alpha能更快地响应不同架构的性能差异。

核心技巧与避坑指南：

• 热身（Warm-up）：在训练初期（例如前10个epoch），先固定架构参数alpha，只训练网络权重W。这有助于超网络先学到一些基础的特征表示，避免alpha在随机初始化的权重下做出错误的判断。
• 交替优化（Bilevel Optimization）：一种更稳定的策略是模仿DARTS，将优化过程视为一个双层问题。在每一步，先在一个minibatch上更新网络权重W，然后在另一个（或同一个）minibatch上，根据当前W下的验证损失来更新架构参数alpha。虽然实现稍复杂，但能缓解alpha和W相互干扰的问题。
• 正则化与DropPath：超网络容易过拟合。除了常用的权重衰减，在可搜索的Cell中应用DropPath（随机丢弃整条路径）至关重要。它能提高超网络的泛化能力，使得评估出的alpha更可靠。DropPath的概率可以线性增加，例如从0逐渐增加到0.2。
• 内存管理：超网络包含了所有可能的操作，非常消耗显存。务必使用梯度检查点（Gradient Checkpointing）技术，它用计算时间换显存空间，能让你在有限的GPU上训练更大的搜索空间。

4.2 架构搜索阶段

当超网络训练收敛后，架构参数alpha就包含了每个选择的重要性信息。搜索阶段就是根据alpha导出最终的最优架构。

导出方法：

1. 确定性子导出：对于每个可搜索的位置（如一条边），选择alpha值最大的那个操作。例如，在某条边上，3x3 Conv的alpha是0.6，5x5 Dilated Conv是0.3，Identity是0.1，那么就选择3x3 Conv作为该边的最终操作。
2. 基于验证集的微调选择：确定性导出可能不是最优的。更稳健的做法是：根据alpha生成多个候选架构（例如top-k个），然后在验证集（或从训练集划出的一部分）上对每个候选架构进行短时间（如5-10个epoch）的重新训练（从超网络继承权重），然后选择验证集性能最好的那个。这个过程被称为“选择后微调”，虽然增加了计算量，但能显著提升最终架构的性能。
3. 考虑硬件约束：如果你有部署需求（如延迟、FLOPs限制），需要在搜索时加入约束。可以在导出时，先根据alpha采样大量架构，快速估算它们的FLOPs或延迟（使用预构建的查找表或轻量级预测器），过滤掉不符合约束的，再从符合条件的架构中按上述方法选择性能最好的。

一个常见的陷阱：直接导出的架构在独立训练时性能下降。这通常是因为超网络训练中存在“马太效应”——强的操作在训练中越来越强，挤压了其他操作的训练机会，导致超网络对其权重估计过高。解决方案是：在超网络训练中加强对较弱操作的探索，例如在采样时引入一定的随机性（epsilon-greedy策略），或者使用更公平的权重共享策略（如FairNAS、SPOS提出的单路径方法）。

5. 从搜索到部署：模型训练、蒸馏与效率优化

搜索得到最终架构只是第一步。这个“新生儿”网络还需要经过精心的培养（训练）和瘦身（优化），才能在实际任务中发挥威力。

5.1 独立训练（Training from Scratch）

将搜索到的确定架构实例化，从头开始训练。这是标准流程，目的是让网络权重针对该特定架构进行充分优化。

• 训练技巧：可以使用在ImageNet上预训练的权重初始化骨干网络（如果骨干是基于常见设计如ResNet变体）。对于完全新颖的架构，则只能随机初始化。此时，数据增强、学习率调度（如Cosine Annealing）、优化器选择（如AdamW）等标准技巧都非常重要。由于架构是专门为任务搜索的，它通常能比手工设计的基线模型更快收敛，达到更高的精度上限。
• 超参数调整：搜索出的架构可能对某些超参数更敏感。建议进行小范围的超参数搜索，如初始学习率、权重衰减系数、批大小等。

5.2 知识蒸馏（Knowledge Distillation）

如果搜索过程使用了代理任务（如小分辨率图像），那么最终架构在大分辨率图像上训练时，可能会丢失一些在搜索阶段学到的“知识”。此外，超网络本身可以看作是一个强大的教师模型。我们可以使用知识蒸馏来提升最终学生模型的性能。

• 离线蒸馏：将训练好的超网络或一个在代理任务上表现优秀的复杂模型作为教师，用它产生的软标签（Soft Targets）来指导最终学生模型的训练。损失函数结合了真实标签的交叉熵损失和与教师输出的KL散度损失。
• 在线蒸馏：在超网络训练阶段，就鼓励不同子架构之间互相学习，让它们产生的特征或预测相互对齐。这能提升超网络的整体质量，从而让搜索出的架构起点更高。

5.3 模型压缩与部署优化

搜索出的架构可能仍包含冗余。为了部署到资源受限的设备，还需进一步优化：

• 通道剪枝（Channel Pruning）：基于重要性评估（如L1范数、BN层缩放因子），剪掉卷积层中不重要的输出通道。
• 量化（Quantization）：将模型权重和激活从32位浮点数（FP32）转换为8位整数（INT8），可以大幅减少模型大小和加速推理。可以使用训练后量化（PTQ）或量化感知训练（QAT）。
• 硬件感知搜索：更高级的做法是，在最初的NAS搜索空间中就引入硬件延迟查找表（Look-Up Table），将目标设备的实际推理延迟作为优化目标之一，直接搜索出既准又快的架构。这需要使用如ProxylessNAS、FBNet等硬件感知NAS方法。

6. 实战常见问题、排查技巧与效果评估

在实际操作中，你会遇到各种各样的问题。下面是我总结的一些典型问题及其排查思路。

6.1 超网络训练不稳定，损失震荡或发散

• 可能原因1：学习率设置不当。架构参数alpha的学习率相对于网络权重W的学习率过大或过小。
  • 排查：监控alpha值的变化。如果alpha变化过于剧烈（很快收敛到0或1），说明其学习率可能太大。如果几乎不变，则可能太小。
  • 解决：尝试调整alpha优化器的学习率，通常设置为W优化器学习率的5-10倍是一个不错的起点。使用学习率warmup。
• 可能原因2：搜索空间太大或存在“死操作”。某些操作（如大卷积核）可能因为难以优化而始终得不到训练，导致其梯度不稳定。
  • 排查：检查每个候选操作被采样到的频率和其对应的alpha值。如果某个操作的alpha始终接近0且很少被采样，它可能就是“死操作”。
  • 解决：简化搜索空间，移除明显低效的操作。或者在训练初期，强制均匀采样所有操作一段时间，给每个操作公平的“起跑线”。
• 可能原因3：未使用DropPath或正则化不足。
  • 解决：确保在可搜索的Cell中应用了DropPath，并适当增加权重衰减。

6.2 搜索出的架构性能不如手工设计的基线模型

• 可能原因1：代理任务与真实任务差异过大。例如，用极低分辨率图像搜索，得到的架构可能缺乏处理高分辨率细节的能力。
  • 解决：尽可能让搜索阶段的设置（如图像分辨率、数据集规模）接近最终任务。如果资源有限，可以尝试使用更小的完整数据集，而不是用大数据的极小部分。
• 可能原因2：架构评估方法有偏。直接根据alpha导出可能不可靠。
  • 解决：务必进行“选择后微调”。即导出几个候选架构，在验证集上快速重训练评估，选择最好的。
• 可能原因3：搜索空间设计有缺陷，遗漏了关键组件。
  • 解决：分析搜索出的架构，看其结构是否合理。例如，在分割任务中，如果搜索出的架构完全没有跳跃连接或特征融合，那说明搜索空间设计可能有问题，需要将必要的连接方式纳入可搜索范围。

6.3 最终模型训练过拟合

• 可能原因：搜索到的架构可能非常贴合搜索/验证集，但对训练数据的泛化能力不足。
  • 解决：加强数据增强。使用更激进的数据增强策略，如MixUp、CutMix、强大的几何与色彩变换。此外，在独立训练阶段也可以使用标签平滑、更强的Dropout/DropPath等正则化技术。

6.4 效果评估指标解读

在密集预测任务中，不能只看一个指标。以语义分割为例：

评估时，需要在验证集上报告mIoU和Pixel Acc，同时在测试集上做最终报告。对于部署，必须将精度（mIoU）与效率（FPS/FLOPs）放在一起对比，绘制帕累托前沿图（Pareto Frontier），清晰展示不同架构在精度-效率权衡上的位置。一个成功的密集预测NAS项目，其搜索出的架构应该在这个帕累托前沿上优于手工设计的基线模型。

最后，我想分享一点个人体会：NAS不是银弹，它更像一个强大的“架构工程师助理”。它无法替代你对任务本质（密集预测需要什么）、数据特性、硬件约束的深刻理解。你的这些理解，最终都体现在搜索空间的设计和搜索目标的制定上。一个精心设计的、融入领域知识的搜索空间，配合高效的搜索策略，才能让NAS发挥出最大价值。刚开始尝试时，不妨从一个小的、可控的搜索空间开始，比如只搜索特征融合部分，而固定一个成熟的骨干网络（如ResNet）。快速验证流程跑通后，再逐步扩大搜索范围。这个过程里，耐心地调试超参数、分析中间结果，比盲目追求复杂的算法更重要。