1. 项目概述当NAS遇上密集预测在计算机视觉领域密集预测任务Dense Prediction Tasks——比如语义分割、深度估计、光流计算、全景分割——一直是模型设计的“硬骨头”。这些任务要求模型对输入图像的每一个像素都做出精确的预测或分类其计算复杂度和对细节信息的需求远非简单的图像分类可比。传统的做法是由领域专家凭借深厚的经验手动设计出像U-Net、DeepLab、HRNet这样的经典网络架构。这个过程耗时费力且严重依赖“炼丹师”的直觉和运气。神经网络架构搜索Neural Architecture Search, NAS的出现一度被视为解放生产力的“自动化炼丹炉”。它通过算法自动在庞大的搜索空间中寻找最优的网络结构在图像分类任务上取得了巨大成功甚至发现了超越人工设计的架构。然而当我们将这个“自动化炼丹炉”直接搬到密集预测任务面前时却发现事情远没有那么简单。搜索空间的设计、评估指标的选取、计算成本的爆炸每一个环节都充满了新的挑战。这个项目就是一次深入探索NAS如何适配密集预测任务并试图攻克其中核心难题的实践记录。我结合了近期的研究和项目经验梳理了从理论到落地的完整链路希望能为同样想在这个交叉领域“掘金”的朋友们提供一份避坑指南。2. 核心挑战与设计思路拆解2.1 密集预测任务对NAS提出的独特要求为什么直接把在ImageNet上搜出来的分类网络拿来做分割效果往往不佳核心在于任务本质的差异。图像分类关注的是全局的、高层的语义信息而密集预测需要同时处理全局上下文和局部细节。这直接导致了几个根本性的设计冲突首先多尺度特征融合成为刚需。一个优秀的密集预测网络必须能有效融合来自骨干网络Backbone浅层的高分辨率、低语义特征利于定位边缘和深层的低分辨率、高语义特征利于识别类别。NAS在分类任务上的搜索空间往往侧重于如何堆叠卷积、设计高效的模块如MBConv, Inverted Residual但对特征金字塔、跳跃连接Skip Connection等融合结构的设计关注不足。其次计算开销的评估标准变了。分类任务通常只看FLOPs或参数量。但在密集预测中尤其是部署到边缘设备时内存带宽Memory Bandwidth和激活值Activation的大小可能成为更关键的瓶颈。因为高分辨率的特征图在层间传递时会消耗巨大的内存读写开销。一个FLOPs很低的模型可能因为频繁在内存中搬运巨大的特征图而实际推理速度极慢。最后搜索目标损失函数需要重新设计。分类用交叉熵损失顶多加个标签平滑。而密集预测任务往往使用复合损失函数例如分割任务常用的交叉熵损失Dice损失的组合。在NAS过程中是直接优化这个最终任务损失还是设计一个与架构强相关的代理损失如网络延迟需要仔细权衡。2.2 主流NAS范式在密集预测中的适配策略面对这些要求我们通常有三种NAS范式可以选择基于强化学习RL的、基于进化算法EA的和基于梯度Gradient-based的。目前基于梯度的可微分架构搜索DARTS及其变种因其相对较高的搜索效率在研究中更受青睐。但在密集预测场景下我们需要对其进行大幅改造。搜索空间的重新定义是第一步。我们不能只搜索骨干网络必须将解码器Decoder或特征融合网络Neck也纳入搜索范围。一个典型的搜索空间可能包含骨干网络单元延续DARTS的风格搜索卷积、可分离卷积、空洞卷积、池化等操作的连接方式。特征融合路径搜索从骨干网络不同阶段stage提取特征图进行融合的方式例如是相加Add还是拼接Concat以及融合前的上采样/下采样方法。多尺度注意力机制搜索是否插入及在何处插入注意力模块如SE、CBAM以增强模型对上下文信息的捕捉能力。一次性One-shotNAS与超网络Supernet几乎是当前处理密集预测NAS的标配。因为从头训练每个候选架构来评估其性能成本是无法承受的。我们需要训练一个包含了所有可能子架构的超网络然后通过权重共享和子网络采样来评估架构性能。这里的关键挑战在于如何设计超网络的训练策略如均匀采样、可微分采样来缓解权重共享带来的评估偏差一个常见的技巧是在训练超网络时对解码器部分也进行路径采样迫使超网络学习到更具泛化性的特征。3. 核心细节解析与实操要点3.1 构建面向密集预测的搜索空间设计搜索空间是NAS成功的基石。一个过于宽泛的空间会让搜索难以收敛而一个过于狭窄的空间则可能找不到最优解。对于密集预测我建议采用分层搜索空间。3.1.1 骨干网络搜索空间这部分可以借鉴成熟的移动端NAS工作如MobileNetV3 EfficientNet但需要调整。例如在较浅的层处理高分辨率特征图应限制那些会急剧扩大通道数的操作以控制内存占用。我们可以将搜索空间定义为一系列有序的“块Block”每个块从一组候选操作中选择候选操作集 O {3x3 深度可分离卷积, 5x5 深度可分离卷积, 3x3 空洞可分离卷积扩张率2, 恒等映射Identity, 零操作Zero}。其中零操作用于实现网络深度的动态调整。每个块还需要搜索扩展率Expansion Ratio和输出通道数但为了简化通常可以固定几个离散选项如 {3, 6} 和 {16, 24, 32, 40}。3.1.2 特征融合网络搜索空间这是密集预测NAS的核心。以经典的U型结构为例我们可以将解码器的每一层定义为一个可搜索的“融合单元”。输入选择对于解码器的第i层其输入可以来自编码器对应的第j层特征跳跃连接以及解码器第i-1层上采样后的特征。我们可以搜索是否使用跳跃连接甚至从多个编码器层中进行选择多跳跃连接。融合操作选择特征融合的方式候选集 F {拼接Concat, 逐元素相加Add, 注意力加权融合Attention-based Fusion}。上采样方法在融合前需要对来自深层或上一解码层的低分辨率特征进行上采样。候选集 U {双线性插值Bilinear Upsample, 转置卷积Transposed Conv, 像素洗牌Pixel Shuffle}。注意拼接操作会显著增加通道数从而增加后续卷积的计算量。在搜索时需要将其计算成本体现在约束条件中否则搜索算法会倾向于生成所有路径都拼接的“胖子”网络导致推理缓慢。3.2 超网络训练与架构评估策略超网络的训练质量直接决定了架构评估的可靠性。一个经典的训练流程如下初始化构建包含所有可能路径的超网络初始化架构参数α决定操作权重和网络权重ω。交替优化步骤一训练网络权重ω从超网络中采样一个批次Batch的子网络。采样策略可以是均匀的也可以根据当前的架构参数α进行概率采样。用目标任务如分割的损失函数L_train来更新这些子网络的权重ω。由于权重共享这次更新会影响到所有共享这些权重的其他子网络。步骤二更新架构参数α固定网络权重ω在另一个独立的验证集上通过前向传播计算当前架构参数α对应的性能如验证集上的mIoU。然后通过梯度下降法更新α目标是最大化验证集性能。在DARTS中这通过双层优化实现。实操心得验证集的选择至关重要。绝对不能使用与训练集有重叠的验证集因为超网络容易过拟合到训练集。最好能从数据集中专门划分出一部分作为“架构验证集”这部分数据不参与网络权重ω的更新只用于评估α。否则搜索过程会倾向于选择那些在训练集上过拟合最快的架构而非泛化能力最好的。另一个常见陷阱是“架构塌缩Architecture Collapse”搜索过程后期架构参数α会倾向于极端的“one-hot”分布即几乎所有权重都集中到某一个操作上通常是参数量最大的操作如5x5卷积。为了缓解这个问题可以在优化α时对其施加L1正则化鼓励稀疏性但避免过度集中或者采用Progressive Shrinking策略在搜索初期允许更复杂的操作组合随着搜索进行逐步剪枝掉权重较低的操作。4. 实操过程与核心环节实现4.1 基于可微分搜索DARTS变种的实现框架这里我以PyTorch环境为例简述一个适配语义分割任务的简化版DARTS实现框架。我们假设搜索空间包含骨干网络单元和简单的融合单元。4.1.1 定义可搜索单元Searchable Cell首先我们需要定义混合操作MixedOp它是所有候选操作的加权和。import torch import torch.nn as nn import torch.nn.functional as F class MixedOp(nn.Module): 混合操作在推理时根据架构参数alpha选择权重最大的操作 def __init__(self, C_in, C_out, stride, ops_list): super().__init__() self._ops nn.ModuleList() for op_name in ops_list: # 实例化候选操作例如3x3 sep conv, 5x5 sep conv, identity, zero等 op build_op(op_name, C_in, C_out, stride) self._ops.append(op) # 架构参数长度等于候选操作数量 self.alpha nn.Parameter(torch.randn(len(ops_list)) * 1e-3) def forward(self, x): # 训练时加权求和 weights F.softmax(self.alpha, dim0) output sum(w * op(x) for w, op in zip(weights, self._ops)) # 推理时argmax选择单一操作 # output self._ops[torch.argmax(self.alpha)](x) return output4.1.2 构建超网络接着我们用这些可搜索单元来搭建一个简化的编码器-解码器超网络。class SuperNet(nn.Module): def __init__(self, num_cells, C_init, num_classes): super().__init__() # 初始的stem层固定结构 self.stem nn.Sequential(nn.Conv2d(3, C_init, 3, padding1), nn.BatchNorm2d(C_init), nn.ReLU()) # 编码器部分多个可搜索的Cell self.encoder_cells nn.ModuleList() C_curr C_init for i in range(num_cells): # 每个Cell的输入输出通道数、步长可根据需要设计 cell SearchableCell(C_curr, C_curr*2 if i%21 else C_curr, stride2 if i%21 else 1) self.encoder_cells.append(cell) C_curr C_curr*2 if i%21 else C_curr # 解码器部分简单的融合单元这里简化为可搜索的上采样卷积 self.decoder_ops nn.ModuleList() for i in reversed(range(num_cells)): # 解码器单元融合来自编码器对应层的特征 dec_op MixedOp(C_curr encoder_channels[i], C_curr//2, stride1, ops_list[conv_3x3, dilated_3x3]) self.decoder_ops.append(dec_op) C_curr C_curr // 2 # 最终预测头 self.head nn.Conv2d(C_curr, num_classes, 1) def forward(self, x, architectureNone): # stem x self.stem(x) encoder_features [] # 编码器前向传播保存中间特征 for cell in self.encoder_cells: x cell(x) encoder_features.append(x) # 解码器前向传播逐步融合特征 for i, dec_op in enumerate(self.decoder_ops): # 获取对应的编码器特征 enc_feat encoder_features[-(i1)] # 上采样当前特征图到与编码器特征相同尺寸 x F.interpolate(x, sizeenc_feat.shape[2:], modebilinear, align_cornersTrue) # 拼接特征 x torch.cat([x, enc_feat], dim1) # 通过可搜索的融合操作 x dec_op(x) # 最终预测 out self.head(x) out F.interpolate(out, scale_factor4, modebilinear, align_cornersTrue) # 上采样回原图尺寸 return out4.1.3 实现交替优化训练循环核心的训练循环伪代码如下# 初始化超网络和架构参数优化器 model SuperNet(...) optimizer_weights torch.optim.SGD(model.parameters(), lr0.025, momentum0.9, weight_decay3e-4) optimizer_alphas torch.optim.Adam(model.arch_parameters(), lr3e-4, betas(0.5, 0.999)) for epoch in range(total_epochs): # 阶段一训练网络权重 (w) model.train() for batch in train_loader: optimizer_weights.zero_grad() # 采样一个子网络在实际中可以通过设置alpha的梯度掩码或随机路径实现 loss model.compute_loss(batch) # 分割损失如CrossEntropy Dice loss.backward() optimizer_weights.step() # 阶段二更新架构参数 (alpha) model.eval() # 注意更新alpha时通常不启用BatchNorm的统计量更新 for batch in val_loader: # 使用独立的架构验证集 optimizer_alphas.zero_grad() # 注意这里需要计算验证集损失关于alpha的梯度 # 在DARTS中这通过一个近似的二阶优化实现计算开销大 # 简化版可以直接在验证集上做前向传播计算损失然后backward这近似为一阶优化 val_loss model.compute_loss(batch) val_loss.backward() optimizer_alphas.step() # 可选应用架构参数的正则化或约束 apply_alpha_regularization(model.arch_parameters())4.2 搜索后的架构导出与从头训练搜索过程结束后我们根据学习到的架构参数α导出最终的网络结构。通常的做法是在每个混合操作处保留权重最大的前k个操作k1或2从而得到一个确定的、精简的网络。关键步骤架构解析遍历超网络中的所有MixedOp根据alpha参数选择Top-k操作并实例化对应的确定操作如3x3深度可分离卷积丢弃其他操作。构建最终网络根据解析出的操作连接关系用标准的PyTorch模块如nn.Conv2d,nn.BatchNorm2d重新搭建一个全新的、确定的网络。这个网络不包含任何权重共享或参数化选择。权重继承可选但推荐将超网络中对应操作的权重拷贝到新网络的对应层中。这可以作为新网络训练的高质量预训练权重显著加速收敛并可能提升最终性能。但要注意操作类型和形状必须完全匹配。从头训练Retraining使用完整的训练数据集对这个确定的网络进行标准的、充分的训练。千万不要直接使用搜索阶段的超网络权重作为最终模型因为权重共享机制下的权重是次优的。5. 常见问题与排查技巧实录在实际操作中你会遇到各种各样的问题。下面是我踩过的一些坑和对应的解决方案。5.1 搜索过程不稳定验证集性能剧烈波动现象在更新架构参数α的阶段验证集上的损失或指标如mIoU像坐过山车一样时高时低无法稳定优化。可能原因与排查验证集太小或与训练集分布不一致检查你的架构验证集是否足够大且有代表性。尝试增大验证集规模或使用K折交叉验证中的一折作为固定的架构验证集。超网络权重ω训练不充分在更新α之前ω必须达到一个相对稳定的状态。如果ω还在剧烈变化那么基于它评估的α梯度就是不可靠的。解决方法是增加每个epoch中训练ω的迭代次数或者先预热Warm-up训练ω多个epoch再开始交替优化α。优化器或学习率设置不当用于α的优化器通常是Adam学习率可能太高。尝试降低lr_alpha例如从3e-4降到1e-4。梯度冲突ω和α的优化目标可能存在冲突。可以尝试更先进的优化策略如PC-DARTS中提出的部分通道连接来减轻优化难度。5.2 搜索出的架构性能不如手工设计的基准模型现象费尽心思搜索出来的网络在相同训练设置下mIoU比经典的U-Net或DeepLabv3还低。可能原因与排查搜索空间设计有缺陷你的搜索空间可能根本没有包含足够好的候选操作。例如如果没把空洞卷积Atrous Conv或空间金字塔池化ASPP纳入候选模型就很难处理大尺度上下文信息。检查搜索空间确保包含了该领域公认有效的模块。代理任务或损失函数不合理如果你为了加速搜索使用了分辨率更低的图片或更少的训练轮次作为代理任务那么这个代理任务可能与最终的全分辨率、全量数据训练任务存在优化目标偏差。确保代理任务尽可能贴近真实任务。架构-权重失配Architecture-Weight Mismatch超网络中共享的权重对于某些架构可能是好的但对另一些可能很差。当导出最终架构时继承的这些权重可能并不适合它。务必进行彻底的重训练Retraining并且可以尝试不同的学习率调度和优化器。过拟合架构验证集搜索算法可能找到了一个在特定验证集上表现极好但泛化能力差的架构。使用数据增强并在最终评估时使用独立的测试集。5.3 搜索耗时过长无法承受现象跑一次搜索实验需要几周甚至几个月无法快速迭代想法。可能原因与排查搜索空间过大这是最主要的原因。评估你的搜索空间大小。如果每个节点有8种操作有10个节点理论组合数就是天文数字。必须进行空间剪枝借鉴已知的高效架构设计先验如MobileNet的倒残差结构设计更小的、结构化的搜索空间。未使用权重共享或One-shot方法确认你使用的是基于超网络的一次性NAS方法而不是像早期RL-NAS那样每个架构都从头训练。代理任务不够“代理”在搜索阶段使用更低分辨率的输入图像如256x256而非512x512、更小的批量大小Batch Size、更少的训练周期Epoch和更小的模型宽度Channel数。这些都能极大降低单次迭代的成本。研究表明在低分辨率下搜索到的架构在高分辨率下微调后通常依然有效。硬件限制NAS本身就是计算密集型的。考虑使用多GPU进行数据并行训练超网络。对于架构搜索部分如果可能将其分配到多个GPU上并行评估不同的子网络尽管在可微分NAS中这较难实现但在进化算法中很常见。5.4 导出架构时出现形状不匹配错误现象在根据α导出最终网络并尝试加载超网络权重时出现Tensor形状不匹配的RuntimeError。排查与解决逐层对比这是最繁琐但最有效的方法。编写脚本打印出超网络中目标MixedOp在选定操作后的权重形状例如selected_op.conv.weight.shape再打印出新网络中对应层的权重形状。99%的错误源于这里。检查跳跃连接在密集预测网络中形状错误经常发生在解码器的特征融合处。确保上采样操作双线性插值、转置卷积的输出尺寸与编码器对应特征的尺寸完全一致。由于舍入误差有时F.interpolate需要明确指定output_size参数而不是仅靠scale_factor。通道数对齐当融合操作选择“拼接Concat”时导出网络的对应层输入通道数会是两部分之和。而超网络中的MixedOp可能是在一个不同的通道数下训练的。确保在导出时你正确计算了拼接后的通道数并初始化了正确形状的新卷积层。个人心得NAS不是“一劳永逸”的银弹而是一个强大的设计空间探索工具。它最大的价值不在于一定能找到“天下第一”的架构而在于它能系统性地探索人类设计师容易忽略的结构组合尤其是在处理像密集预测这样多目标、多约束的复杂任务时。在项目初期与其投入巨大资源做全自动搜索不如先构建一个合理的搜索空间然后用NAS跑出几个有潜力的候选架构再结合人工经验进行分析和微调。这种“人机协同”的模式往往比完全依赖算法更高效、更可靠。最后一定要把计算成本约束如延迟、内存明确地加入到搜索目标中否则搜出来的模型很可能无法落地。