1. 为什么值得花时间啃透DARTS如果你正在做神经网络架构搜索NAS方向的研究或者工程落地大概率绕不开DARTS这篇工作。全称Differentiable Architecture Search核心思路一句话就能说清楚把离散的架构选择松弛成连续的、可微的权重然后用梯度下降同时优化网络权重和架构参数。这个想法在2018年提出来的时候直接炸了NAS这个圈子因为它把原本需要几千GPU天的搜索过程压缩到了几个GPU天。我自己第一次读这篇论文的时候感觉公式推导很顺代码应该也不难。结果真正把官方代码拉下来跑通、再自己改结构的时候才发现坑远比想象中多。从搜索阶段的显存爆炸到离散化之后的性能崩塌再到评估阶段的复现差异每一步都有大量论文里不会写的细节。这篇文章就是把我从读论文到跑代码、再到自己改结构这一整条链路上积累的东西全部摊开讲包括核心原理的拆解、代码关键模块的逐行分析、实操中遇到的典型问题和排查思路。适合的读者有三类一是刚接触NAS、想通过DARTS入门的研究生或者算法工程师二是已经跑过官方代码但结果对不上、想搞清楚哪里出了问题的同学三是想把DARTS的思想迁移到自己业务场景里的从业者。不管你是哪一类我都会尽量把“为什么这么做”讲透而不是只贴一段代码让你抄。2. DARTS的核心设计思路拆解2.1 从离散搜索到连续松弛这个弯是怎么拐过来的传统的NAS方法比如基于强化学习或者进化算法的方案本质上是在一个离散的搜索空间里做黑盒优化。搜索空间里每个候选架构都是一个独立的点你需要完整训练它、评估它才能知道它好不好。这就导致搜索成本极高因为每次评估都意味着一次完整的训练过程。DARTS的关键洞察在于我不需要把每个候选架构当成独立的点我可以把搜索空间表示成一张有向无环图每条边上的操作是一个混合操作混合的权重就是架构参数。这样一来架构搜索就变成了对这些连续权重的优化问题可以直接用梯度下降来解。具体来说假设一个单元cell里有若干节点节点之间的边上有若干候选操作比如3x3卷积、5x5卷积、跳连、池化等。DARTS的做法是在每条边上放一个softmax加权的混合操作混合操作 sum over o of (exp(alpha_o) / sum over o of exp(alpha_o)) * o(x)其中alpha就是架构参数。搜索的过程就是同时优化网络权重w和架构参数alpha。训练完成后每条边保留softmax权重最大的那个操作得到一个离散的架构。这个松弛的好处是显而易见的搜索过程变成了可微的可以用标准的反向传播来优化效率提升了好几个数量级。但代价也很明显松弛之后的连续空间和离散空间之间存在gap搜索时最优的连续架构离散化之后不一定还是最优的。这个问题在后面会详细展开。2.2 双层优化为什么不能简单地把alpha和w一起训DARTS的优化目标是一个双层优化问题bilevel optimization。外层优化架构参数alpha目标是让在验证集上的损失最小内层优化网络权重w目标是让在训练集上的损失最小。用公式表示就是min over alpha of L_val(w*(alpha), alpha) s.t. w*(alpha) argmin over w of L_train(w, alpha)为什么不能直接把alpha和w放在一起用训练集损失来优化因为那样会导致架构参数过拟合训练集搜索出来的架构在验证集上表现很差。双层优化的设计就是为了让架构参数在验证集上做选择保证泛化性。但双层优化有个麻烦w*(alpha)本身是内层优化的结果要对alpha求导需要用到隐函数定理或者近似。DARTS采用的是一种一阶近似的方法通过一步梯度下降来近似w*(alpha)然后对alpha求导。这个近似在论文里有详细的推导代码里也有对应的实现。实际操作中这个近似会带来一些问题。比如搜索阶段的batch size不能太小否则梯度的方差会很大导致alpha的优化不稳定。另外训练集和验证集的划分也很关键如果验证集太小或者分布和训练集差异太大搜索出来的架构可能会偏向于某些特定的操作。2.3 搜索空间设计cell结构的选择与权衡DARTS的搜索空间是基于cell的。一个cell是一个有向无环图包含若干个节点和边。DARTS定义了两种cellnormal cell和reduction cell。Normal cell保持特征图的空间分辨率不变reduction cell将空间分辨率减半。每个cell有两个输入节点分别是前两层的输出和一个输出节点。中间的节点通过边连接每条边上的操作从预定义的候选操作集合中选择。候选操作通常包括3x3深度可分离卷积、5x5深度可分离卷积、3x3空洞卷积、5x5空洞卷积、3x3最大池化、3x3平均池化、identity跳连、zero无连接。这个设计的好处是搜索空间相对紧凑搜索出来的cell可以堆叠成任意深度的网络。但也有一些限制比如cell的结构是固定的只有边上的操作是可搜索的节点的连接方式也是预定义的不能搜索。这些限制在一定程度上约束了搜索空间的表达能力但换来了搜索效率的大幅提升。我在实际使用中发现候选操作集合的选择对搜索结果影响很大。如果池化操作太多搜索出来的架构会倾向于大量使用池化导致模型容量不足如果卷积操作太多搜索又会偏向于堆叠卷积参数量爆炸。通常的做法是保持一个相对平衡的候选集合并根据具体任务做微调。3. 代码关键模块逐行解析3.1 搜索空间的构建从代码角度看cell是怎么定义的官方代码里搜索空间的构建主要在model_search.py这个文件里。核心类是Cell它接收一个steps参数表示cell里有多少个中间节点和一个multiplier参数表示输出节点由多少个中间节点拼接而成。每个cell的构建过程大致是这样的首先定义两个输入节点前两层的输出然后对于每个中间节点从前面的所有节点中选择两个作为输入在这两条边上分别放置混合操作。混合操作的权重就是架构参数alpha。代码里有一个关键的数据结构self._ops它是一个ModuleList每个元素是一个MixedOp对应一条边上的混合操作。MixedOp内部维护了一个ModuleList包含所有候选操作以及一个alpha参数用于计算softmax权重。class MixedOp(nn.Module): def __init__(self, C, stride): super(MixedOp, self).__init__() self._ops nn.ModuleList() for primitive in PRIMITIVES: op OPS[primitive](C, stride, False) self._ops.append(op) def forward(self, x, weights): return sum(w * op(x) for w, op in zip(weights, self._ops))这段代码看起来很简单但有几个细节值得注意。第一weights是通过对alpha做softmax得到的这个softmax是在forward里计算的而不是在__init__里。第二sum这里用的是Python内置的sum而不是torch.sum这意味着每个操作的输出是逐个相加的显存占用会比较高。第三OPS是一个字典把操作名映射到具体的nn.Module这个字典的定义在文件开头。在实际使用中我发现MixedOp的显存占用是一个大问题。因为每条边上都要计算所有候选操作的输出然后加权求和所以显存占用是候选操作数量的倍数。如果候选操作有8个那显存占用就是普通网络的8倍左右。这也是为什么DARTS搜索阶段需要比较大的显存。3.2 架构参数的初始化与更新alpha的优化细节架构参数alpha的初始化方式对搜索结果有一定影响。官方代码里alpha是用torch.nn.Parameter定义的初始化为0。这意味着初始时所有操作的softmax权重是均匀的每条边上的混合操作是所有候选操作的等权平均。alpha的更新是通过优化器来完成的。官方代码里用的是torch.optim.Adam学习率设置为3e-4权重衰减设置为1e-3。这个学习率比网络权重的学习率通常0.025小很多原因是架构参数需要更精细的调整学习率太大会导致搜索过程震荡。在train_search.py里有一个关键的步骤是计算架构参数的梯度。由于双层优化的存在这个梯度不能直接用验证集损失对alpha求导得到而是需要先对网络权重做一步虚拟更新然后再计算梯度。代码里对应的逻辑是# 计算虚拟更新后的网络权重 unrolled_model self._compute_unrolled_model(input, target, eta, network_optimizer) # 计算架构参数的梯度 unrolled_loss self._loss(unrolled_model, self._steps, input_valid, target_valid) unrolled_loss.backward()这里的_compute_unrolled_model就是做一步虚拟更新eta是虚拟更新的学习率。这个步骤是DARTS的核心也是计算量最大的部分因为它需要保留计算图显存占用会进一步增加。我在实际使用中发现如果显存不够可以把unrolled设置为False改用一阶近似。这样虽然搜索精度会有所下降但显存占用会大幅降低。对于显存有限的场景这是一个很实用的折中方案。3.3 离散化与架构导出从连续权重到最终网络搜索完成后需要把连续权重离散化得到最终的架构。这个过程在model_search.py的genotype方法里实现。核心逻辑是对于每条边选择softmax权重最大的操作对于每个中间节点选择权重最大的两条边作为输入。def genotype(self): def _parse(weights): gene [] n 2 start 0 for i in range(self._steps): end start n W weights[start:end].copy() edges sorted(range(i 2), keylambda x: -max(W[x][k] for k in range(len(W[x])) if k ! PRIMITIVES.index(none)))[:2] for j in edges: k_best None for k in range(len(W[j])): if k ! PRIMITIVES.index(none): if k_best is None or W[j][k] W[j][k_best]: k_best k gene.append((PRIMITIVES[k_best], j)) start end n 1 return gene # ...这段代码有几个细节需要注意。第一none操作被排除在外也就是说如果某条边的none操作权重最大这条边会被忽略。第二每个中间节点选择两条输入边这是DARTS的默认设置。第三gene是一个列表每个元素是一个(操作名, 输入节点索引)的元组。离散化之后通常会重新训练一个网络来评估架构的性能。这个重新训练的过程和普通网络训练类似但有一些细节需要注意比如学习率调度、数据增强、正则化等。官方代码里用的是标准的训练流程但在实际使用中这些超参数对最终性能影响很大。4. 实操过程中的关键环节与避坑指南4.1 环境配置与数据准备从零开始跑通官方代码先把环境搭起来。官方代码基于PyTorch建议用1.0以上的版本但不要用太新的版本因为有些API可能不兼容。Python版本建议3.6到3.8之间。CUDA版本根据你的显卡来选但要注意PyTorch版本和CUDA版本的对应关系。数据准备方面官方代码支持CIFAR-10和ImageNet。CIFAR-10可以直接用torchvision.datasets.CIFAR10下载ImageNet需要自己准备。数据增强方面官方代码用了标准的随机裁剪和随机翻转另外还用了cutout。Cutout的实现很简单就是在图像上随机挖一个16x16的方块把它置零。这个操作对防止过拟合很有效。class Cutout(object): def __init__(self, length): self.length length def __call__(self, img): 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 - self.length // 2, 0, h) y2 np.clip(y self.length // 2, 0, h) x1 np.clip(x - self.length // 2, 0, w) x2 np.clip(x self.length // 2, 0, w) mask[y1: y2, x1: x2] 0. img torch.from_numpy(img.numpy() * mask) return img这段代码看起来简单但有一个坑img.numpy()要求img是CPU上的tensor如果数据加载器返回的是CUDA tensor这里会报错。所以cutout要在数据加载阶段做而不是在训练循环里做。4.2 搜索阶段的参数设置与显存优化搜索阶段的参数设置很关键。官方代码里搜索阶段训练50个epochbatch size设置为64网络权重的学习率是0.025架构参数的学习率是3e-4。这些参数在CIFAR-10上表现不错但换到其他数据集可能需要调整。显存优化是搜索阶段的一个大问题。因为每条边上都要计算所有候选操作的输出显存占用是普通网络的数倍。如果显存不够可以尝试以下几个方案减小batch size。这是最直接的方法但会影响搜索稳定性。减少候选操作数量。比如去掉一些不常用的操作只保留核心的几个。使用一阶近似。把unrolled设置为False可以大幅降低显存占用。使用混合精度训练。PyTorch的amp模块可以自动做混合精度显存占用可以减少30%到50%。我自己的经验是如果显存只有8GB建议用一阶近似加混合精度batch size设置为32左右。如果显存有16GB以上可以用二阶近似batch size设置为64。4.3 评估阶段的重新训练与性能复现搜索完成后会得到一个离散的架构。这个架构需要重新训练来评估性能。官方代码里评估阶段训练600个epochbatch size设置为96学习率用余弦退火调度初始学习率0.025。评估阶段有几个坑需要注意。第一数据增强的强度要够否则容易过拟合。除了标准的随机裁剪和翻转还可以加cutout、mixup、cutmix等。第二学习率调度很关键余弦退火比阶梯下降效果更好。第三权重衰减要设置合理太大或太小都会影响性能。我在复现官方结果的时候发现评估阶段的性能波动很大。同样的架构不同的随机种子最终精度可能差0.5%到1%。所以如果要对比不同架构的性能最好用多个随机种子跑取平均值。另外官方代码里的评估阶段用的是CIFAR-10的完整训练集但搜索阶段用的是训练集的50%作为训练集另外50%作为验证集。这个划分方式在评估阶段要改回来否则会浪费一半的训练数据。5. 常见问题与排查技巧实录5.1 搜索阶段loss不下降或者震荡怎么办搜索阶段loss不下降或者震荡通常有几个原因。第一学习率设置不当。架构参数的学习率如果太大会导致搜索过程震荡如果太小搜索会停滞。建议从3e-4开始试如果震荡就降到1e-4如果停滞就升到1e-3。第二数据划分不合理。如果验证集太小或者分布和训练集差异太大架构参数的梯度会很不稳定。建议验证集至少占训练集的30%到50%。第三候选操作集合不合理。如果候选操作里有太多相似的操作比如同时有3x3卷积和5x5卷积搜索过程可能会在这两个操作之间反复横跳。建议候选操作集合保持多样性但不要有太多冗余。第四batch size太小。搜索阶段的batch size如果太小梯度的方差会很大导致搜索不稳定。建议batch size至少设置为32如果显存允许设置为64更好。5.2 离散化之后性能崩塌的原因分析离散化之后性能崩塌是DARTS最被诟病的问题之一。原因主要有两个一是连续松弛和离散选择之间的gap二是搜索阶段的过拟合。连续松弛的gap是指搜索阶段最优的连续架构离散化之后不一定还是最优的。因为连续架构是所有候选操作的加权平均而离散架构只保留权重最大的那个操作。如果权重分布比较均匀离散化之后的信息损失会很大。搜索阶段的过拟合是指架构参数在搜索过程中可能过拟合了验证集。因为搜索阶段只训练50个epoch网络权重还没有充分收敛架构参数可能选择了一些在短期训练中表现好、但长期训练中表现差的操作。缓解这个问题的方法有几个。第一增加搜索阶段的epoch数让网络权重更充分地收敛。第二在搜索阶段加入正则化比如dropout或者权重衰减。第三在离散化之后对架构做微调比如调整某些边的操作看看性能是否提升。第四使用一些改进的DARTS变体比如PC-DARTS、P-DARTS等这些方法在搜索稳定性和离散化性能上都有改进。5.3 显存不足的排查与优化方案显存不足是DARTS实操中最常见的问题。排查思路是这样的首先看搜索阶段的batch size是不是太大如果是先减小batch size。然后看候选操作数量是不是太多如果是减少候选操作。然后看是不是用了二阶近似如果是改成一阶近似。最后看是不是可以用混合精度训练。如果以上方法都不行还可以考虑梯度累积。梯度累积的思路是用小的batch size计算梯度累积多次之后再更新参数。这样可以在不增加显存占用的情况下模拟大的batch size。PyTorch里实现梯度累积很简单就是在backward之后不立即step而是累积几次之后再step。for i, (input, target) in enumerate(train_loader): output model(input) loss criterion(output, target) loss loss / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这段代码里accumulation_steps是累积次数。需要注意的是loss要除以accumulation_steps否则梯度会放大。5.4 常见问题速查表问题现象可能原因排查方法解决方案搜索阶段loss震荡架构参数学习率太大打印alpha的梯度范数降低学习率到1e-4搜索阶段loss不下降学习率太小或数据划分不合理检查验证集大小和分布增大学习率或调整数据划分离散化后性能崩塌连续松弛gap或搜索过拟合对比搜索和评估阶段的精度增加搜索epoch或使用改进变体显存不足batch size太大或候选操作太多用nvidia-smi监控显存减小batch size或减少候选操作评估阶段精度波动大随机种子影响用多个随机种子跑取平均值或增加训练epoch训练速度慢数据加载瓶颈或模型太大用profiler分析增加数据加载线程或减小模型6. 从DARTS出发的扩展与改进方向6.1 搜索效率的进一步优化DARTS虽然比强化学习的方法快很多但搜索阶段仍然需要几个GPU天。后续有很多工作试图进一步优化搜索效率。比如PC-DARTS通过部分通道连接来减少显存占用和计算量P-DARTS通过逐步增加网络深度来稳定搜索过程Fair-DARTS通过公平比较不同操作来避免搜索偏差。这些改进的核心思路大同小异要么减少搜索阶段的计算量要么提高搜索的稳定性。PC-DARTS的部分通道连接是一个很实用的技巧它把每个操作的输入通道分成两部分一部分用于搜索一部分用于保留信息。这样可以在不损失性能的情况下把显存占用降低一半左右。6.2 搜索空间的扩展与迁移DARTS的搜索空间最初是为图像分类设计的后来被扩展到其他任务比如语义分割、目标检测、语音识别等。扩展搜索空间的关键是定义合适的候选操作和cell结构。比如在语义分割里需要加入空洞卷积和上采样操作在目标检测里需要加入多尺度特征融合的操作。迁移到其他数据集时需要注意数据集的规模和复杂度。如果数据集比较小搜索阶段容易过拟合需要加强正则化如果数据集比较大搜索阶段需要更多的epoch来充分收敛。6.3 与其他技术的结合DARTS可以和很多其他技术结合比如知识蒸馏、量化、剪枝等。知识蒸馏可以用一个大模型来指导搜索过程让搜索出来的架构更好地模仿大模型的行为。量化可以在搜索阶段就考虑量化误差搜索出对量化更友好的架构。剪枝可以在搜索完成后对架构做进一步的压缩。这些结合方式的核心思想是把DARTS作为一个架构搜索的工具和其他技术协同使用而不是孤立地看待。在实际业务场景里往往需要综合考虑精度、速度、模型大小等多个指标单一的技术很难满足所有需求。7. 我个人在实际操作中的几点体会跑DARTS这段时间最大的感受是论文里的公式很优雅但代码里的细节很磨人。很多在论文里一笔带过的东西在实际操作中都是坑。比如数据划分的比例、学习率的调度、候选操作的选择这些在论文里可能只有一句话但在代码里需要反复调试。另外一个体会是不要迷信官方代码的结果。官方代码在CIFAR-10上能跑到97%以上的精度但换到其他数据集或者换个随机种子结果可能差很多。所以复现的时候一定要多跑几次取平均值不要被单次结果迷惑。最后DARTS的思想比具体的实现更重要。连续松弛、双层优化、cell结构这些思想可以迁移到很多其他问题里。即使你不做NAS了解这些思想对理解其他优化问题也有帮助。