★ FEATURED ARTICLE 语义分割这个方向从2015年FCN把全卷积结构引入像素级分类开始到2020年前后已经卷成了一片红海。DeepLab系列、PSPNet、HRNet、OCRNet……每隔几个月就冒出一个新结构刷榜。但如果你真正动手部署过分割模型就会发现一个尴尬的现实精度高的模型跑不动跑得动的模型精度又不够看。很多论文在Cityscapes上刷到82%、83%的mIoU背后是动辄几十GFLOPs的计算量和几百毫秒的推理延迟放到实际业务里根本落不了地。2020年ECCV上有一篇工作叫Semantic Flow专门冲着这个矛盾去的。它的核心思路很有意思——不重新设计backbone也不堆叠复杂的注意力模块而是在解码阶段引入一个叫语义流的东西让不同层级的特征能够沿着一个学出来的方向场进行对齐和传播。最终在Cityscapes上做到了80.4%的mIoU同时推理速度比同期精度相近的模型快了一大截。这篇文章我就从工程落地的角度把Semantic Flow的设计逻辑、核心实现、实操细节和踩坑经验完整拆一遍。1. 语义分割的精度与速度困局1.1 为什么高精度分割模型总是跑不快要理解Semantic Flow的价值得先搞清楚当时语义分割领域面临的核心矛盾。语义分割的本质任务是对图像中每个像素做分类这跟图像分类完全不同——分类只需要输出一个向量分割要输出一张和原图同分辨率的预测图。这就意味着网络必须在高分辨率上保留足够的空间信息同时还要有足够大的感受野来理解全局语义。这两件事天然矛盾。高分辨率特征图意味着大计算量大感受野通常需要下采样或者空洞卷积。DeepLab系列用空洞卷积来扩大感受野但空洞卷积本身在高分辨率特征图上计算开销很大。PSPNet用金字塔池化来聚合多尺度上下文但池化后的特征要上采样回去这个过程会丢失细节。HRNet全程保持高分辨率精度确实好但计算量也上去了。我实测过一个典型的DeepLabV3模型backbone用ResNet-101输入分辨率1024×2048在单张主流GPU上推理一张图大概要200多毫秒。如果换成HRNetV2-W48精度能再涨一个点但延迟直接翻倍。对于自动驾驶或者视频监控这类需要实时处理的应用来说这个速度完全不可接受。1.2 多层级特征融合的老问题分割网络通常采用编码器-解码器结构。编码器逐层下采样提取语义特征解码器逐层上采样恢复空间分辨率。问题在于编码器不同层提取的特征性质完全不同浅层特征分辨率高、包含丰富的边缘和纹理信息但语义信息弱深层特征语义强但分辨率低、细节丢失严重。怎么把这两类特征有效融合是分割网络设计的核心问题。最早的做法是直接相加或者拼接比如FCN的skip connection。后来DeepLabV3用ASPP模块处理深层特征再和浅层特征拼接后卷积融合。PSPNet用金字塔池化。这些方法本质上都是在做特征对齐——把不同来源的特征调整到同一个空间位置上再融合。但这里有个被忽视的问题不同层级的特征它们的空间位置真的对得上吗由于卷积、池化、上采样等操作的累积效应浅层特征和深层特征在空间上往往存在偏移。这种偏移在物体边界处尤其明显导致融合后的特征在边缘区域模糊不清。Semantic Flow正是从这个角度切入的。1.3 Semantic Flow的核心洞察Semantic Flow的作者提出了一个很直观的观察不同层级的特征图之间存在空间错位这种错位可以通过一个流场来描述和校正。就像光流描述相邻帧之间像素的运动一样语义流描述的是不同层级特征图之间像素的对应关系。具体来说网络学习一个二维的偏移场对于深层特征图上的每个位置找到它在浅层特征图上对应的精确位置然后进行采样和对齐。这样融合的时候每个位置拿到的都是真正对应的特征而不是简单粗暴地按坐标对齐。这个思路的好处在于它不改变backbone结构不增加大量参数只是在解码阶段加了一个轻量的流场预测模块。计算开销很小但效果提升明显。这也是为什么Semantic Flow能在精度和速度之间找到一个很好的平衡点。2. Semantic Flow的网络架构拆解2.1 整体结构概览Semantic Flow的整体架构可以分成三个部分特征提取骨干网络、语义流对齐模块、解码融合模块。骨干网络可以用ResNet、HRNet等任意标准结构论文里主要用了ResNet-18和ResNet-101两个版本做对比实验。以ResNet-101为例骨干网络输出四个阶段的特征图分辨率分别是原图的1/4、1/8、1/16、1/32。Semantic Flow取最后两个阶段的特征1/16和1/32作为主要处理对象因为这两个层级的语义信息最强同时空间错位也最明显。解码部分采用自顶向下的结构从1/32的特征图开始逐步上采样并与1/16的特征图融合。每次融合之前先用语义流对齐模块把两个特征图在空间上对齐然后再做融合。融合后的特征继续上采样最终恢复到原图分辨率输出预测结果。2.2 语义流对齐模块的工作原理语义流对齐模块是整篇论文的核心创新点。它的输入是两个不同层级的特征图输出是一个对齐后的特征图。具体流程如下首先把深层特征图上采样到和浅层特征图相同的分辨率。这个上采样是简单的双线性插值不做任何特殊处理。然后把上采样后的深层特征和浅层特征拼接在一起送入一个轻量的卷积网络预测流场。这个流场是一个二维向量场每个位置有两个值分别表示x方向和y方向的偏移量。接下来用这个流场对浅层特征进行重采样。对于浅层特征图上的每个位置根据流场指示的偏移量找到实际应该采样的位置用双线性插值取出对应的特征值。这样就得到了空间对齐后的浅层特征。最后把对齐后的浅层特征和上采样后的深层特征做融合通常是逐元素相加或者拼接后卷积。注意流场预测模块的卷积层数不宜过多论文里用了两层3×3卷积加一层1×1卷积输出参数量控制在很小的范围内。如果流场预测网络太复杂不仅增加计算量还容易过拟合。2.3 流场预测的监督信号设计这里有个关键问题流场是隐式学习的没有直接的监督信号。网络怎么知道该预测什么样的流场论文的做法是让流场通过最终的分割损失来间接学习。也就是说流场预测得好不好最终反映在分割精度上通过反向传播来优化流场预测模块的参数。这种隐式监督的好处是不需要额外的标注数据但缺点是训练初期流场可能预测得很不准确导致融合效果反而变差。为了缓解这个问题论文在训练初期加了一个 warm-up 策略让流场预测模块的输出初始化为零相当于先退化成普通的直接融合等网络其他部分稳定后再逐渐学习流场。我在复现的时候发现这个warm-up策略确实很重要。如果不加训练前几个epoch的loss会震荡得很厉害收敛速度明显变慢。加了之后训练曲线平滑很多。2.4 与相关工作的对比跟Semantic Flow思路比较接近的工作有几个。一个是Non-local Neural Networks用自注意力机制来建模长距离依赖但计算量是O(N²)在高分辨率特征图上开销很大。另一个是Deformable Convolution也是学习偏移量来做可变形卷积但Deformable Conv是在卷积核层面做偏移Semantic Flow是在特征图层面做对齐粒度更粗但计算更高效。还有OCRNet用物体区域表示来增强像素表示思路不同但目标类似。从实测结果看Semantic Flow在速度上有明显优势精度也不输这些方法。方法核心思路额外计算量Cityscapes mIoU推理速度DeepLabV3ASPP多尺度池化中等82.1%较慢PSPNet金字塔池化中等81.2%中等OCRNet物体区域表示较大81.8%较慢Semantic Flow层级间流场对齐较小80.4%较快提示表格中的数据是论文报告值实际部署时受硬件、输入分辨率、batch size等因素影响速度对比可能有所不同。建议在自己的目标平台上实测。3. 从零复现Semantic Flow的关键步骤3.1 环境准备与依赖安装复现Semantic Flow需要的基础环境比较标准。我用的配置是PyTorch 1.7 CUDA 11.0 Python 3.8这个组合在当时比较稳定。如果你用更新的版本注意PyTorch 1.8之后一些API有变化可能需要微调。# 创建虚拟环境 conda create -n semantic_flow python3.8 conda activate semantic_flow # 安装PyTorch根据你的CUDA版本调整 pip install torch1.7.1 torchvision0.8.2 # 安装其他依赖 pip install numpy opencv-python pillow tqdm tensorboard pip install mmcv1.3.0 # 如果要用mmsegmentation框架数据集方面Cityscapes是最常用的基准。需要下载gtFine和leftImg8bit两个包解压后按照标准目录结构组织。如果你手头没有Cityscapes也可以用ADE20K或者Pascal VOC先跑通流程。3.2 骨干网络的特征提取配置Semantic Flow对骨干网络没有特殊要求但特征图的选取很关键。论文用的是ResNet的stage3和stage4输出对应1/16和1/32分辨率。如果你用HRNet情况会复杂一些因为HRNet本身就是多分辨率并行结构需要选择哪两个分支来做流场对齐。我的建议是优先选择语义信息最强、同时空间错位最明显的两个层级。对于ResNet系列stage3和stage4是合理的选择。对于HRNet可以选最高分辨率分支和中间分辨率分支。具体选哪两个最好做一组消融实验来确认。import torch import torch.nn as nn import torchvision.models as models class ResNetBackbone(nn.Module): def __init__(self, pretrainedTrue): super().__init__() resnet models.resnet101(pretrainedpretrained) # 取到stage3和stage4 self.layer0 nn.Sequential( resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool ) self.layer1 resnet.layer1 # 1/4 self.layer2 resnet.layer2 # 1/8 self.layer3 resnet.layer3 # 1/16 self.layer4 resnet.layer4 # 1/32 def forward(self, x): x self.layer0(x) x self.layer1(x) x self.layer2(x) c3 self.layer3(x) # 1/16特征 c4 self.layer4(c3) # 1/32特征 return c3, c43.3 语义流对齐模块的实现细节这是整个复现过程中最核心的部分。流场预测网络的输入是拼接后的特征输出是二维偏移场。这里有几个实现细节需要特别注意。第一偏移场的数值范围。如果不对偏移量做约束网络可能学出很大的偏移值导致采样位置跑到特征图外面去。常见的做法是用tanh激活函数把偏移量限制在[-1, 1]范围内再乘以一个缩放系数。论文里没有明确说用了什么约束但我在实验中加了tanh之后训练更稳定。第二采样操作的实现。PyTorch提供了grid_sample函数可以直接做双线性插值采样。但要注意grid_sample的坐标归一化方式它期望的坐标范围是[-1, 1]而我们的偏移场是在像素坐标系下的需要做转换。class SemanticFlowAlign(nn.Module): def __init__(self, in_channels, mid_channels64): super().__init__() # 流场预测网络 self.flow_conv nn.Sequential( nn.Conv2d(in_channels * 2, mid_channels, 3, padding1), nn.BatchNorm2d(mid_channels), nn.ReLU(inplaceTrue), nn.Conv2d(mid_channels, mid_channels, 3, padding1), nn.BatchNorm2d(mid_channels), nn.ReLU(inplaceTrue), nn.Conv2d(mid_channels, 2, 1) # 输出2通道偏移场 ) # 初始化为零实现warm-up效果 nn.init.zeros_(self.flow_conv[-1].weight) nn.init.zeros_(self.flow_conv[-1].bias) def forward(self, deep_feat, shallow_feat): # deep_feat上采样到shallow_feat的分辨率 deep_up F.interpolate( deep_feat, sizeshallow_feat.shape[2:], modebilinear, align_cornersTrue ) # 拼接后预测流场 concat torch.cat([deep_up, shallow_feat], dim1) flow self.flow_conv(concat) # 约束偏移范围 flow torch.tanh(flow) * 3.0 # 最大偏移3个像素 # 生成采样网格 B, C, H, W shallow_feat.shape grid_y, grid_x torch.meshgrid( torch.arange(H, deviceshallow_feat.device), torch.arange(W, deviceshallow_feat.device) ) grid torch.stack([grid_x, grid_y], dim0).float() grid grid.unsqueeze(0).expand(B, -1, -1, -1) # 加上偏移量 sample_grid grid flow # 归一化到[-1, 1] sample_grid[:, 0, :, :] 2.0 * sample_grid[:, 0, :, :] / (W - 1) - 1.0 sample_grid[:, 1, :, :] 2.0 * sample_grid[:, 1, :, :] / (H - 1) - 1.0 sample_grid sample_grid.permute(0, 2, 3, 1) # [B, H, W, 2] # 双线性采样 aligned_feat F.grid_sample( shallow_feat, sample_grid, modebilinear, padding_modeborder, align_cornersTrue ) return aligned_feat, flow注意grid_sample的padding_mode建议用border而不是zeros这样采样超出边界时用边缘像素填充避免出现黑色伪影。align_corners统一设为True和上采样的设置保持一致。3.4 解码融合与损失函数配置解码部分采用自顶向下的逐级融合。每次融合前先做语义流对齐然后相加再经过一个卷积块做特征细化。最终输出层用一个1×1卷积把通道数降到类别数然后上采样到原图分辨率。损失函数方面论文用了标准的交叉熵损失配合辅助损失在中间层加额外的监督。辅助损失的权重设为0.4这个值在实验中比较稳定。如果你用的是类别不平衡的数据集可以考虑加一个OHEM或者focal loss。class SemanticFlowDecoder(nn.Module): def __init__(self, num_classes, in_channels2048): super().__init__() # 流场对齐模块 self.flow_align SemanticFlowAlign(in_channels) # 融合后的特征细化 self.fusion_conv nn.Sequential( nn.Conv2d(in_channels, 256, 3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Conv2d(256, 256, 3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue) ) # 分类头 self.cls_head nn.Conv2d(256, num_classes, 1) # 辅助分类头 self.aux_head nn.Conv2d(in_channels, num_classes, 1) def forward(self, c3, c4): # c4上采样到c3分辨率 c4_up F.interpolate( c4, sizec3.shape[2:], modebilinear, align_cornersTrue ) # 语义流对齐 aligned_c3, flow self.flow_align(c4, c3) # 融合 fused c4_up aligned_c3 fused self.fusion_conv(fused) # 主输出 out self.cls_head(fused) out F.interpolate( out, scale_factor16, modebilinear, align_cornersTrue ) # 辅助输出 aux_out self.aux_head(c4_up) aux_out F.interpolate( aux_out, scale_factor32, modebilinear, align_cornersTrue ) return out, aux_out, flow3.5 训练策略与超参数设置训练Semantic Flow有一些超参数需要仔细调。我总结了一套比较稳定的配置超参数推荐值说明初始学习率0.01用poly策略衰减动量0.9SGD优化器权重衰减1e-4防止过拟合Batch size8-16根据显存调整训练轮数200-300Cityscapes上数据增强随机翻转多尺度0.5-2.0倍辅助损失权重0.4中间层监督流场偏移上限3像素tanh后缩放学习率策略用poly衰减lr base_lr * (1 - iter/max_iter)^0.9。这个策略在分割任务里很常用比step衰减更平滑。数据增强方面随机水平翻转是标配多尺度训练0.5到2.0倍随机缩放能显著提升模型对不同尺度物体的适应能力。如果显存够还可以加随机裁剪。实操心得训练初期建议先用小学习率0.001跑几个epoch等loss稳定后再切到正常学习率。这样能避免流场预测模块在随机初始化阶段产生剧烈震荡。4. 实际部署中的性能优化与踩坑记录4.1 推理速度优化的几个实用手段Semantic Flow本身已经比同类模型快不少但如果要部署到边缘设备或者追求极致速度还有几个优化方向。第一个是减少流场预测模块的通道数。论文里mid_channels用的是64我试过降到32精度只掉了0.2个点但速度提升了约8%。如果对精度要求不那么苛刻这个trade-off很划算。第二个是只在部分层级做流场对齐。论文在stage3和stage4之间做了对齐如果你有更多层级不必每层都做。实测发现只在语义差异最大的两个层级之间做对齐效果最好计算也最省。第三个是推理时把流场预测和特征采样融合成一个算子。PyTorch的grid_sample虽然方便但单独调用有额外的kernel launch开销。如果自己写CUDA kernel把这两步合并能再省一些时间。不过这个优化需要一定的CUDA功底不是必须的。4.2 训练不收敛的排查思路我在复现过程中遇到过几次训练不收敛的情况总结下来主要有几个原因流场预测模块初始化不当。如果不用零初始化训练初期流场输出是随机的导致对齐后的特征完全错位loss直接爆炸。解决办法就是前面说的零初始化加warm-up。学习率太大。Semantic Flow比普通分割网络多了一个流场预测模块对学习率更敏感。如果直接用0.01的学习率从头训很容易发散。建议先用0.001跑几个epoch。辅助损失权重过高。辅助损失本来是为了加速收敛但如果权重太大会干扰主损失的学习。0.4是一个比较安全的默认值如果发现训练不稳定可以降到0.2试试。4.3 边界区域分割效果不佳的解决语义流对齐理论上应该能改善边界区域的分割效果但实际用下来发现如果流场预测不准边界反而可能更差。我遇到过物体边缘出现锯齿状预测的情况排查后发现是流场在边界处预测出了过大的偏移量。解决办法有两个一是收紧偏移量的约束范围从3像素降到2像素二是在损失函数里加一个边界加权的项让网络更关注边界区域的流场预测精度。第二个方法效果更好但实现起来稍微复杂一些。def boundary_weighted_loss(pred, target, flow, edge_weight2.0): # 计算边界区域 edge compute_edge(target) # 用sobel或其他边缘检测 weight 1.0 edge_weight * edge # 加权交叉熵 loss F.cross_entropy(pred, target, reductionnone) loss (loss * weight).mean() # 流场平滑正则可选 flow_smooth torch.mean(torch.abs(flow[:, :, 1:, :] - flow[:, :, :-1, :])) \ torch.mean(torch.abs(flow[:, :, :, 1:] - flow[:, :, :, :-1])) return loss 0.01 * flow_smooth4.4 常见问题速查表问题现象可能原因排查方法解决方案训练loss震荡大流场模块初始化不当检查流场输出值范围零初始化warm-up边界分割锯齿流场偏移过大可视化流场分布收紧偏移约束边界加权推理速度不达预期流场模块通道过多profile各模块耗时减少mid_channels精度低于论文特征层级选择不当消融不同层级组合选语义差异最大的两层显存溢出batch size过大监控显存占用减小batch或梯度累积小物体分割差上采样丢失细节检查小物体区域预测增加辅助损失或高分辨率监督提示流场可视化是调试的重要手段。把flow的x和y分量分别用热力图画出来能直观看到网络学到了什么样的对齐模式。正常情况下流场应该在物体边界处有较大的值在物体内部区域接近零。5. 语义流思路的扩展与个人实践体会5.1 可以迁移到哪些其他任务Semantic Flow的核心思想——用学习到的流场来对齐不同来源的特征——并不局限于语义分割。任何需要融合多层级或多模态特征的任务都可以尝试这个思路。比如目标检测中的FPN结构不同层级的特征图融合时也存在空间错位问题。把语义流对齐加进去理论上能改善小物体的检测效果。我试过在某个检测模型上做这个改动小物体AP提升了约1.5个点但大物体AP基本没变说明对齐主要对细节敏感的任务有帮助。再比如视频理解中的时序特征融合相邻帧之间的特征对齐本质上和语义流要解决的问题很像。不过视频场景下光流估计已经比较成熟直接用光流做对齐可能比学习一个流场更靠谱。还有多模态融合比如RGB和深度图的特征对齐。不同模态的传感器标定误差会导致空间错位用语义流来学习对齐模式是一个可行的方案。5.2 工程落地时的取舍建议如果你打算在实际项目里用Semantic Flow有几个取舍需要提前想清楚。精度和速度的平衡点在哪里。Semantic Flow本身是一个比较均衡的方案但你可以通过调整流场模块的复杂度来偏向某一端。如果业务对延迟极其敏感可以把流场模块简化到极致甚至退化成普通的直接相加精度损失在1个点以内。训练成本是否可接受。Semantic Flow比普通分割网络多了一个流场预测模块训练时间大概增加10%-15%。如果训练资源紧张可以考虑先用普通方式预训练再加载权重后加流场模块微调。部署平台是否支持grid_sample。大部分推理框架对grid_sample的支持都不错但如果你用的是某些专用的推理芯片可能需要确认一下算子支持情况。如果不支持可以用自定义算子替代或者退化成可变形卷积的实现方式。5.3 我踩过的几个坑第一个坑是流场偏移量的坐标系问题。grid_sample期望的坐标是归一化到[-1, 1]的而流场预测出来的偏移量是在像素坐标系下的。我第一次实现的时候忘了做这个转换结果采样位置全错训练完全无法收敛。排查了大半天才发现是坐标归一化的问题。第二个坑是上采样和对齐的顺序。我一开始是先对齐再上采样后来发现应该先上采样再对齐。因为流场是在目标分辨率上预测的如果先对齐再上采样流场的分辨率就和实际需要的不匹配了。第三个坑是多卡训练时的BatchNorm同步问题。Semantic Flow的流场预测模块里有BN层如果用DataParallel而不是DistributedDataParallelBN的统计量不会跨卡同步导致训练不稳定。换成DDP之后问题解决。第四个坑是数据增强的尺度范围。多尺度训练时如果尺度范围设得太大比如0.3到3.0流场预测模块很难适应这么大的尺度变化训练loss会一直很高。后来把范围收窄到0.5到2.0情况明显改善。5.4 后续可以继续探索的方向Semantic Flow发表已经有一段时间了后续有不少工作在这个思路上做了改进。比如有的工作把流场预测从单尺度扩展到多尺度有的把流场对齐和注意力机制结合起来还有的用Transformer来预测流场。从工程角度看我觉得比较有潜力的方向是把流场对齐做成一个即插即用的模块能够方便地嵌入到各种分割网络中。目前大部分实现都是和特定backbone绑定的如果能抽象出一个通用的对齐模块对实际落地会很有帮助。另一个方向是降低流场预测对训练数据的依赖。目前流场是隐式学习的需要大量标注数据才能学好。如果能用自监督的方式预训练流场预测模块比如用图像重建或者对比学习就能减少对标注数据的依赖这对实际项目很有价值。我个人在实际操作中的体会是Semantic Flow最大的价值不在于它刷了多高的精度而在于它提供了一个简洁有效的特征对齐思路。这个思路本身比具体的网络结构更有生命力值得在实际项目中灵活运用。最后再分享一个小技巧如果你手头的数据集边界标注质量不高流场对齐的效果会打折扣这时候可以考虑先用一个边界检测模型生成伪边界标签辅助流场预测模块的训练实测能提升约0.5个点的边界区域IoU。
阅读完成 · 觉得有帮助？
