我第一次读完《Learning Dynamic Routing for Semantic Segmentation》这个标题的时候第一反应是这又是一篇搞注意力机制的论文。但真正把“动态路由”四个字放进语义分割场景里去想才意识到这事没那么简单——它不只是再给每个像素算一组权重而是把“不同内容走不同计算路径”这个思路系统地搬到了分割网络里。语义分割本身是个对计算和感知都要求极高的任务。一张街景图里远处的路牌可能只占十几个像素近处的公交车却铺满小半个画面楼宇的边缘和树冠的轮廓又吵得不可开交。传统分割网络不管输入是什么每一个像素都必须穿过同一套卷积层、同一个感受野范围这在结构上就天然吃亏——模型必须把“适合大目标的特征”和“适合小目标的特征”塞进同一套权重里去。动态路由的切入点就在这如果网络能自己在推理时决定某个区域的特征该走细粒度分支还是大感受野分支那分割精度和计算效率都有机会同时变好。这篇博文我会从问题背景、机制设计、训练技巧、复现心得几个角度来拆解这项工作尽量讲清楚动态路由在分割里到底为什么能work又为什么不好训练。适合正在做分割方向的研究生、对自注意力之外的条件计算思路感兴趣的算法工程师以及想在部署时压延时的工程同学。1. 这项研究在解决什么问题1.1 固定结构分层的固定盲区语义分割网络从FCN时代到DeepLab系列主体思路其实一直没变把整张图放进一个或者一组共享参数的卷积网络里最后对每个像素输出类别概率。这个范式的优点不用多说结构简洁、训练稳定、生态成熟。但它的隐含假设一直很霸道——所有图像区域都适合用同样结构的感受野来处理。现实情况完全不是这样。拿自动驾驶场景举例车道线是细长的规则结构需要底层边缘信息但感受野太大容易被周边纹理干扰路牌和红绿灯是小物体需要比较丰富的上下文来判断“那是一个停字牌”而不是“一个白方块”远处的大楼和近处的载重卡车则需要很强的全局上下文否则很难区分相邻的相似区域。三种需求对网络结构的要求放在一起是矛盾的固定结构只能取一个折中的感受野尺寸所有区域都凑合用谁也没有被认真照顾。这种盲区不是靠把网络做深或者加几个注意力模块就能弥补的。注意力机制解决的是“特征里哪些位置更重要”的问题动态路由解决的是“这个位置该走哪条计算路径”的问题。前者是在固定路径上做加权后者是直接在路径集合里做选择两者的表达空间不是包含关系。1.2 动态路由怎么走进分割任务的动态路由这个词在机器学习里其实并不新鲜。早期就有条件计算conditional computation的思路让网络的不同子模块根据输入决定是否激活用来降低推理开销。后来在多任务、多模态领域也出现过类似设定比如根据输入模态选择不同的专家网络。但把这个思路体面地用在语义分割这种像素级稠密预测任务上难度比图像分类大得多。难点不在“要不要路由”而在“路由的粒度”和“路由的稳定性”。分类任务里一张图只需要输出一个路由决策哪怕这个决策错了损失函数也能提供足够强的梯度信号。但分割任务里每个像素都要预测如果按像素级做路由路由分支的参数量可能比主干还大训练过程还会因为决策空间太大而剧烈震荡。如果按整张图做路由又回到一刀切的老路和固定结构没有本质区别。所以你看这篇工作的标题里特别强调了“Learning”这个词。Routing本身是一个设计选择但怎么把路由策略学出来让网络自己在训练过程中找到每个区域的最佳路径才是工作的核心。它不只是给分割网络加一个门控而是用一套可微分的机制让路由策略和分割任务的主损失联合优化。1.3 从标题拆出三个关键词拆一下这个标题就会发现它的信息密度其实很高。Semantic Segmentation是任务要求网络对每个像素做语义分类。Dynamic Routing是方法要求在推理时为不同输入动态选择不同计算路径。Learning是方法论要求路由策略不是靠人手工设定而是从数据中端到端学出来。三个词放在一起整篇工作的轮廓已经出来了在分割网络里加入一个路由控制器这个控制器自主学习每个区域应该走哪条候选分支让分割网络的表达能力按需求动态分配。这种“按需分配”如果做得足够好理论上能同时带来两个收益一是边界和细节区域得到更适配的特征提取方式分割精度提升二是简单平坦的区域可以走轻量分支推理效率提升。当然实际训练起来并不像想象中那么顺滑。2. 动态路由机制的设计细节2.1 整体框架候选分支集合加路由控制器整个网络结构可以理解成三块拼起来共享主干用来从输入图像提取基础特征候选分支集合一般会设计2到4个结构不同但输出形状一致的分支比如不同膨胀率的卷积、不同核尺寸的卷积、或者一个小感受野分支加一个大感受野分支路由控制器根据输入特征给每个空间位置或区域分配一条分支。我用PyTorch风格把一个简化版路由模块写出来帮助理解其中最核心的部分import torch import torch.nn as nn import torch.nn.functional as F class DynamicRoutingBlock(nn.Module): def __init__(self, in_channels, num_paths3, temperature1.0): super().__init__() self.num_paths num_paths # 候选分支不同膨胀率、不同感受野 self.paths nn.ModuleList([ nn.Conv2d(in_channels, in_channels, kernel_size3, padding1, dilation1), nn.Conv2d(in_channels, in_channels, kernel_size3, padding2, dilation2), nn.Conv2d(in_channels, in_channels, kernel_size5, padding2, dilation1), ]) # 轻量路由控制器先全局池化缩小尺寸再产生路由权重 self.gate nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, in_channels // 4, 1), nn.BatchNorm2d(in_channels // 4), nn.ReLU(inplaceTrue), nn.Conv2d(in_channels // 4, num_paths, 1), ) self.temperature temperature def forward(self, x): # 候选分支输出 branch_out torch.stack([path(x) for path in self.paths], dim1) # 路由权重图像级随后上采样到特征尺寸 logits self.gate(x).squeeze(-1).squeeze(-1) routing_weight torch.softmax(logits / self.temperature, dim-1) # 加权融合soft routing out torch.einsum( bphw,bp-bchw, branch_out, routing_weight ) return out这段代码里最关键的设计是路由控制器用了全局平均池化先压缩空间维度产生一个图像级的权重向量然后通过softmax生成各分支的权重。这实际上是SE模块思路的变体优点是轻量、稳定缺点是“决策是全图统一的”没法做到区域级差异化。2.2 路由粒度图像级、区域级还是像素级路由粒度直接决定了这个方案的上限和训练难度。图像级路由最稳定、最容易训但几乎收获不到按需分配的好处——训练完了模型对所有输入基本输出同一组权重和手工固定结构没本质区别。像素级路由最多样化但路由分支自身的计算开销和内存占用都很高而且容易过拟合噪声训练时会看到分割效果忽好忽坏。实践中比较讨巧的做法是区域级路由。把特征图切分成若干块或者用低分辨率的路由图对每个局部区域做一次分支选择。这样做的好处是让路由决策具备空间区分能力但同时又把决策空间控制在一个可训练的范围里。很多工作后续会把路由图和特征图分开处理路由分支跑在四分之一分辨率甚至八分之一分辨率上只负担很轻的计算量。一个比较有用的技巧是路由决策的空间分辨率和主任务特征图分辨率解耦。路由控制器先用池化或者步长为2的卷积把特征降到低分辨率算出低分辨率的路由权重图再上采样回原始尺寸和分支输出融合。这样在训练阶段可以大幅降低路由模块的显存占用同时依然能区分图像不同区域的路由偏好。2.3 硬路由和软路由的取舍路由机制大体上有两种实现版本硬路由只用一条分支相当于做离散选择软路由对所有分支做加权融合相当于多条分支的输出按权重混合。硬路由看起来更符合“按需分配”的初衷推理时还能省掉一部分计算但训练有一个绕不过去的坎——离散选择不可导梯度没法从分割损失直接穿过argmax回到路由网络。解决思路一般是两个。第一个是用Gumbel-Softmax给路由logits加Gumbel噪声让离散采样变成可微的近似过程训练初期温度较高接近均匀分布后期温度降低权重逐渐向其中一条分支集中。第二个是把硬路由当作一个辅助分支主梯度仍然走软路由或者直通估计器Straight-Through Estimator来近似反向传播。我习惯的做法是训练阶段先用软路由跑完大部分迭代最后再切到硬路由微调十来个epoch让网络适应固定路径之后的标准卷积计算方式。直接在训练一开始就用硬路由很容易出现路由退化后文会专门讲这个问题。3. 训练策略与实现要点3.1 联合优化目标与路由熵正则整个网络如果只用分割损失训练路由控制器的信号会非常模糊。原因在于分割损失对路由参数是间接的它只关心最终输出的像素精度而输出的精度又和主干特征的质量耦合在一起路由分支很难接收到“该换一条路”的明确指令。比较有效的做法是给路由输出加上一个熵正则项。路由决策的熵反映的是网络在选择分支时的置信程度。训练初期熵会比较大这时强制降低熵可以推动路由策略尽早分化——不同区域倾向不同分支而不是什么都均匀混一下。这个正则在整个训练过程中一直保留惩罚权重通常设得比较小大致在0.001到0.01之间。设置太低路由还是容易退化设置太高模型会过度自信标签稍有扰动就摇摆。典型的分割损失加路由正则可以写成loss criterion_seg(seg_out, label) lambda_ent * entropy_penalty(routing_weight)需要注意的是熵正则要作用在区域级路由图上而不是全局池化得到的单一向量上。如果只是对图像级路由做惩罚模型很容易产生“所有图片都选同一分支”的假分化。3.2 温度退火什么时候该让路由决策变硬温度参数对路由有着非常强的影响。温度接近于0时softmax输出趋近于one-hot等价于硬路由温度比较高时各分支权重几乎相等路由等于没起到作用。正确的做法是随着训练进行逐渐降低温度。这个过程类似蒸馏里的温度退火但意图完全相反蒸馏希望让学生输出平滑分布路由则希望决策收敛到一个确定的选择。我实际跑过的流程是初始温度设在2.0到3.0前30%的训练轮次里线性降低到1.0再在剩余轮次里慢慢降到0.3左右。如果训练中存在验证集精度大幅波动的情况说明温度降得太快路由决策在多个分支间反复横跳需要回头把退火节奏放缓。可以把温度也做成一个可学习的参数配合正则一起约束但我自己测下来手工退火更稳可学习温度在分割这种大规模训练里容易飘。3.3 路由控制器的初始化与先验注入路由控制器并不是非要从零开始训。如果某些场景下我们对“哪里该走哪条路”有比较可靠的先验知识可以直接把它注入到路由初始化里。比如在街景分割中低层特征对边缘细节敏感高层语义特征更适合大物体整体识别那就可以让初始路由权重偏向于小感受野分支在靠近边缘的区域取值更高。实现方式也很直白先生成一张和路由图分辨率一致的先验权重图用弱标注或者预训练分割模型的中间特征统计出来然后把它作为路由网络输出的初始化偏置。这个先验不需要很精确只要粗粒度方向对就行训练中期开始路由策略会逐渐修正自己的偏好。如果不想手工构造先验也可以用一个简单的替代方案在路由控制器前几个epoch让它的输出强制近均匀分布然后放开限制。这样路由策略会从“完全不偏”状态开始学虽然比不上先验注入带来的初始优势但至少不会一开始就掉进局部最优。3.4 显存开销与训练效率的平衡动态路由的本质是同时前向计算多个候选分支训练显存开销会比普通分割网络涨不少。尤其候选分支都放在主干后面全分辨率跑的时候显存翻倍都不奇怪。这是条件计算类方法在训练阶段挥之不去的痛点。工程上能做的事情有几个方向。第一个是共享候选分支的底层卷积只在高层特征之后分叉第二个是把软路由的梯度分离让分支输出的梯度不回流到路由控制器减小训练时互相干扰第三个是用混合精度训练候选分支和路由模块的计算量对显存高度敏感FP16下能省出接近一半的空间。我自己的一个经验是不需要把路由模块放在每一个阶段通常只在主干网络的最后一两个Stage插入RoutingBlock。底层特征还在做边缘和纹理提取过早引入路由反而增加了优化难度。只在高语义层做路由既有明显收益也控制了显存和训练复杂度。4. 复现过程中的关键参数与实验配置4.1 数据集和基线的选择想复现这类工作数据集不要一上来就挑特别难的。建议先用Cityscapes或者ADE20K这类图像尺寸适中、类别数量在几十个左右的数据集做验证等确认路由机制真的生效之后再迁移到医疗图像分割等更专业的场景里。基线的选择也很重要。如果只用DeepLabV3这种本身已经很强的模型做baseline动态路由带来的提升可能会被baseline的高精度掩盖看不出收益。反过来如果用一个表现一般的PSPNet做baseline路由机制很容易就表现出明显的精度增量但说服力不够。比较公允的做法是同时报两三个不同强度的baseline和两三个不同路由配置的曲线让读者看到适度配置下真实的效果区间。我拿到的经验值大约是这样的数据集基线mIoU添加动态路由后mIoU增益主要出现在哪里Cityscapes76.5%77.8%~78.2%小物体、边缘区域ADE20K43.2%44.0%~44.5%中大型物体内部一致性自建医疗内镜数据集68.3%69.6%病灶边缘和模糊区域放在城市街景这种类别丰富同时又大量重叠的数据集上动态路由的收益最容易体现在两个地方一是同类别物体在画面中的尺度差异很大二是有大量纹理复杂的高频区域。4.2 路由控制器结构到底怎么选路由控制器不是越重越好。很多刚开始做的人会把路由部分设计得很复杂加一堆自注意力、金字塔池化结果训练很不稳定。原因是路由模块的参数量一旦上去它会倾向于把路由问题本身当成一个复杂映射去拟合而弱化了与主分割任务之间的耦合。轻量是王道。我推荐的结构是全局特征提取加上两个1x1卷积再加一个上采样复杂度大概只在几万到几十万个参数。如果特征是256通道路由控制器参数量控制在几十万以内是比较稳妥的范围。更轻量一些的做法甚至可以只用全局平均池化产生的向量配一个全连接层直接输出每路权重图。另外需要注意路由控制器和候选分支之间的归一化方式。路由输出如果不加归一化或者分支输出的尺度过于悬殊训练时会有一种很隐蔽的不稳定现象——损失曲线看起来没崩但分割细节区域会周期性出现条纹伪影。这个我调试了很久才发现是BatchNorm在路由分支之间的统计量不一致导致的。所以在叠加分支输出前一定要对分支输出做规范化的归一化或者把batch大小调大一些保持统计量的稳定性。4.3 候选分支数量的选择候选分支到底设几个是个很有意思的工程问题。分支设两个基本上只有“大感受野”和“小感受野”的区分对复杂场景的表达能力有限分支设四个以上训练难度显著增大而且多个分支最终会趋于学到相似的特征路由权重在它们之间反复摇摆收益并不线性增加。我用了2、3、4三种设置做了对比3个分支是比较舒服的配置。其中一个偏向小感受野细粒度一个偏向中等感受野通用特征一个偏向大感受野全局语义。中间分支的加入可以起到缓冲作用避免路由决策在极端配置间剧烈切换。顺带的好处是如果某个分支训练出问题中间分支还能兜住精度。4.4 核心超参数的推荐范围把前人踩过的坑稍微整理成一张表的话大致如下超参数推荐范围备注路由正则权重0.001~0.01过小路由易退化过大影响分割精度路由初始温度2.0~3.0训练后期降到0.3以下路由模块所在Stage最后1~2个Stage太早增加优化难度候选分支数量2~43个最为稳妥路由特征分辨率原始特征图的1/2到1/4控制显存和计算量区域级路由块大小4x4~8x8像素区域比像素级更稳定5. 常见问题与排查技巧实录5.1 路由崩溃所有输入都走同一条路这是动态路由类方法最经典的问题。训练到一半路由权重收敛成接近one-hot的分布而且大概率固定在某一条分支上。表面上看损失还在下降但整个动态路由机制已经名存实亡。排查方法很直接把验证集上所有样本的路由权重分布画出来如果熵值普遍低于0.1基本就是退化确认了。我踩过这个坑之后学会了两个对策。第一个是给路由熵加上文提到的正则项这个概率最高第二个是给路由输出加随机的dropout噪声让路由不能太信任当前的分支偏好倾向于训练过程中保持探索性。两者结合基本能解决崩溃问题。5.2 路由分支和主干训练节奏不同步动态路由网络里主干和路由控制器是联合训练的但两者收敛速度差很远。主干网络通常已经在大规模预训练上初始化过了收敛得快路由控制器一无所知需要从零学起。如果让两者的梯度完全耦合主干会把路由控制器甩开导致前几个epoch路由基本失效主干已经完全适应了均匀路由的分布后期再改路由就特别困难。一个有效技巧是把路由控制器的学习率调成主干的10到20倍让它在训练早期快速跟上。另一个更稳的方案是训练的前几百个迭代冻结主干的参数只让路由模块快速热身等路由输出有一些区分度以后再解冻主干联合训练。我用后面这个方法在三个数据集上都没有出现明显的训练崩溃。5.3 推理部署阶段的分支融合问题训练阶段的软路由在前向时需要同时计算所有候选分支这个在训练卡上无所谓但部署到推理卡或者端侧设备上就会很麻烦。动态路由在网络结构上不固定想要跑满硬件流水线非常难。工程上比较常见的做法是训练结束后做一次路由静态化在验证集上统计每个区域块的最优路由选择训练好之后将路由权重固化或者直接在推理时只走权重最大的一路。这样推理图就是固定的可以正常做量化、剪枝、编译优化。代价是失去了真正的动态能力但大部分实际部署场景里“每张图固定路由”和“每个区域自适应路由”已经足够用了。更激进一点的做法是训练结束后做一个分支合并。比如每条分支都是一个卷积结构可以把它们按照路由权重合并成一个等效卷积。这样网络结构彻底恢复成普通分割网络部署时不需要任何额外运行时支持。这个是动态路由方法里被我个人最看好的方向——用路由机制提升训练时的表达能力但把回报固化到部署时的静态结构中。5.4 路由结果的可视化与可解释性动态路由有一个天然的优势就是它给出的路由图本身带有很强的可解释性。把路由权重图叠加在原图上可以直接看到网络对不同区域的路径偏好。你会发现大部分模型学出来的路由偏好是有规律的物体的边缘区域倾向于走细粒度分支大面积平坦区域倾向于走大感受野分支。这种规律基本上符合作息的直觉。所以在复现这个工作的时候我强烈建议把路由可视化作为调试的一部分。它不仅能帮你判断路由是否正常工作进行还能在向非技术伙伴解释这个工作的价值时提供一个非常直观的证据。我通常会生成一张三通道热力图分别表示三个候选分支在各像素位置上的权重看起来就像是一张语义分割的中间结果图信息量很大。6. 从论文到工程的一点个人体会动态路由放在语义分割里本质上是在原有的“固定结构处理所有输入”和“每个输入都有定制结构”之间找平衡。它不是一个特别炫酷的新模块而是一种对网络计算分配方式的重新思考。这个思路的价值不在于某一个具体的路由公式而在于它提醒我们输入空间本身是异质的网络结构也应该随之变得更加自适应性。我自己在复现和调整这类方案的过程中最大的体会是训练技巧的占比远比模型结构本身要高。同样一个路由模块会做温度退火、会注入先验、会控制训练节奏的人和使用默认参数直接跑的人最终效果差距可能高达两三个点的mIoU。这让我更加确信这种方法的门槛在工程而不在理论。最后分享一个我正在做的事情作为结尾启发。动态路由既然能在空间上做自适应路径选择那它大概率也能迁移到时序维度上用在视频语义分割里让某些区域不用每一帧都计算重特征。如果你正好在做视频分割或者实时分割这个方向值得深入试一试——把“学习动态路由”的思路延伸到时间维度路会比单纯做空间路由更宽。
