跨域小样本语义分割：PATnet方法解析与CDFSS基准复现指南

做少样本语义分割的朋友对ECCV2022这篇PATnet应该不陌生。标题里的几个词——Cross-Domain Few-Shot Semantic Segmentation、新基准CDFSS、复现教程——基本把这两年大家最头疼的问题都串起来了标注样本少、训练域和测试域还不一致。PATnet做的最核心的一件事就是给跨域小样本语义分割提供了一套新方法加一套相对完善的评测体系而且论文作者把复现链路也整理得比较清楚。这篇论文分享我尽量不写废话直接把任务背景、方法套路、基准设计和复现要点摊开讲清楚把实际跑代码时容易踩的坑也一起放进去给后面想入坑CDFSS的同学省点时间。这个方向并不是纯学术刷点。举个实际场景你做自动驾驶语义分割手里只有一小批夜间或雨天的标注图模型在白天数据上训练得再好切到夜间往往还是崩。再比如遥感影像分割公开数据集大多是自然场景预训练权重直接迁移到卫星图上效果也很惨。跨域小样本分割要解决的就是这种“少量目标域样本 跨域分布差异”叠加下的快速适应问题属于少样本学习、域适应和语义分割三个方向的交叉点。1. CDFSS到底解决什么问题1.1 从通用分割到小样本分割先说基础设定。常规语义分割任务要求每个类别有大量像素级标注训练一个模型去预测每个像素的类别。这个流程在数据充足时很成熟但标注成本极高尤其是医学影像、遥感影像这种需要专家参与的领域一张图的标注可能比模型训练还贵。小样本分割Few-Shot Semantic Segmentation就是为了解决标注不足的问题给定少量带标注的支持图像support set模型需要在查询图像query set上分割出与支持图像相同类别的目标。通常采用episode的训练方式每个episode包含一个支持集和一个查询集支持集中每个类别只有1到5个标注样本。模型要学会的不是记住某个类别的外观而是学会“如何利用少量标注样本完成分割”这个通用能力。经典的PASCAL-5i、COCO-20i这类基准就是把PASCAL VOC或COCO的类别划分成若干组训练时只用其中一部分类别测试时用未见的类别。这种方式能验证模型对“新类别”的泛化能力但它有一个隐含假设训练和测试数据来自同一个数据分布。实际部署时这个假设往往不成立。1.2 域偏移给少样本分割带来的致命影响当训练集源域和测试集目标域分布不一致时模型的退化速度非常快。比如说在PASCAL VOC自然图像上训练的模型拿到DeepGlobe卫星图上测试mIoU可能直接掉一半甚至更多。原因不难理解特征提取器学习到的颜色、纹理、边缘等低级特征在不同域中差异太大而小样本条件下支持样本又太少不足以对特征分布做有效修正模型很容易把目标域中的背景噪声当成有效特征。这里要区分两个容易混淆的设置域适应Domain Adaptation和跨域小样本分割Cross-Domain Few-Shot Semantic Segmentation。域适应通常假设目标域有大量无标注数据可以通过熵最小化、对抗学习等方式利用这些数据而跨域小样本通常只给目标域提供极少量的标注支持样本比如每个类别1到5个标注甚至连无标注目标域数据都不一定充分。所以在少样本和跨域两个条件叠加时问题比单独任何一个都难。少量支持样本本身可能不能覆盖目标域的真实类内变化而源域预训练模型又带有强烈的源域偏好。如果直接把通用的少样本分割方法拿过来用支持原型和查询特征之间的域差距会导致注意力权重偏移分割掩码常常出现整块误判。1.3 PATnet和新基准在其中的位置PATnet这个工作就是冲着上面那组矛盾来的。它想在同一套框架里同时解决两个问题怎么让支持样本在跨域情况下仍然有效以及怎么让查询特征充分利用支持信息又不被域差异干扰。更直接地说它用了一种原型对齐加Transformer交互的设计把支持原型当作可学习的锚点在注意力计算中逐步把查询特征拉向一个域无关的语义空间。同时在基准层面这篇论文把CDFSS做成了一个相对规范的评测任务。之前这个方向最大的问题是大家各跑各的有人用GTA5到Cityscapes有人用PASCAL到医学影像数据划分、类别对齐、episode采样方式都不一样论文之间没法公平比较。这次提出的CDFSS基准统一了数据组合和评估协议对后续研究来说价值不亚于方法本身。2. PATnet方法拆解跨域自适应怎么落地2.1 整体架构思路原型、注意力、对齐PATnet的完整名称如果按我复现时的理解来解读它强调的是三个关键动作Prototype原型提取、Alignment域对齐、Transformer全局交互。整个框架不是简单地把已有模块堆在一起而是围绕“跨域条件下原型仍然可信”这一个核心目标来设计的。一个典型的小样本分割模型会有这几步用共享特征提取器编码支持图和查询图对支持特征做masked pooling得到每个类别的原型向量然后计算查询特征和原型之间的相似度生成分割预测。跨域场景下这么做有两个问题支持图和查询图的特征分布不在一个空间里直接用余弦相似度或像素级注意力匹配结果会被域偏差带偏另一个问题是支持样本太少单一原型向量表达力有限类内多样性根本体现不出来。PATnet的应对策略可以分成三条线第一条线是用Transformer结构扩大感受野让查询特征在全局上下文中理解语义而不是只看局部像素第二条线是设计对齐机制在特征层面把支持域和查询域的信息相互校准第三条线是改进原型的使用方式不只用单一向量而是结合注意力矩阵做更细粒度的匹配。2.2 核心模块逐一说明从代码实现的角度拆开看PATnet里有几个模块值得细讲。特征提取器部分作者选用了ResNet系列和Swin Transformer作为backbone去掉最后的全连接分类头保留多尺度特征。这里有个小细节主干网络的BatchNorm参数通常被冻结因为小样本episode内的batch非常小只有支持图和查询图各一张或几张BatchNorm统计量不稳定反而会带来噪声。很多复现版本直接沿用这个处理效果也比较稳定。接下来是原型构造模块。不同于常见方法里用masked average pooling直接得到一个类别原型PATnet会先生成多尺度支持特征通过一个轻量卷积层融合再得到原型。这个多尺度融合的收益在于不同域的图像分辨率差异很大卫星图里一个目标可能只占几十个像素自然图里同一类目标可能占上千像素只在单一尺度上提原型信息损失很严重。然后是Transformer交互模块这里是最核心的部分。它以支持原型作为“key/value”的一部分查询特征作为“query”通过多头注意力计算全局关联。值得注意的是它不是简单做一次attention而是多层堆叠并且在每一层都引入了对齐约束。这就让支持信息和查询信息不是一次性完成匹配而是逐层迭代校准。我在阅读代码时发现这个模块输出的不只是分割预测还会输出一个对齐后的查询特征这个特征会送回到解码器做精细分割。最后就是域对齐损失。PATnet在训练时引入了一种对齐损失函数本质上是约束支持特征和查询特征在经过Transformer交互后的分布尽量一致。可以用对抗方式做也可以用MMD或对比学习方式做。从原论文的实验来看这个损失项的权重不能设太大否则模型会为了“对齐”而牺牲分割精度训练过程也不太稳定后面复现部分我会给一个参考值。2.3 与之前方法相比的差异点在PATnet之前跨域小样本分割也有一些探索。一类做法是在源域训练时加入域增强比如随机颜色扰动、风格迁移通过增加数据多样性来提升模型对目标域的鲁棒性。这种方式简单有效但没有从根本上解决域偏移问题遇到域间差异特别大的情况比如自然图像到医学图像效果有限。另一类做法是直接在目标域上微调利用少量支持样本更新整个网络。问题是目标域标注太少全量微调特别容易过拟合而且会遗忘源域学到的基础分割能力。PATnet和这两类都不一样它不依赖大幅度的数据增强也不靠全量微调而是把域对齐做进了注意力计算过程中。支持原型通过Transformer不断和查询特征交互每一层都在修改特征表达相当于一个模型内的隐式适配不需要额外的适配网络。从复现的角度看这种做法工程上更好落地推理时不需要区分源域和目标域是一个端到端的模型。3. 新基准设计CDFSS评测体系3.1 为什么旧评测体系不够用跨域小样本分割此前没有一个公认的评测标准。常见做法是拿通用的FSS基准比如在PASCAL-5i上训练然后换一个目标域测试但目标域选哪个、每个episode的采样方式、类别如何对齐不同论文各有各的搞法。这就导致同一方法在不同论文里的指标差距很大根本没法判断哪个模块真正有效。更麻烦的是很多旧基准存在类别重叠问题。如果在PASCAL上训练时见过了“汽车”这个类测试时目标域里也有“汽车”那么模型记住类别外观也能取得不错效果但这并不能证明它具备跨域少样本泛化能力。真正合理的评测应该保证训练阶段和测试阶段类别完全不重叠同时源域和目标域的图像风格有明显差异。CDFSS基准针对这些问题做了几个关键设计。首先是限定训练类别和测试类别互斥其次是采用统一的域组合定义最后是把支持集数量设置为1-shot和5-shot两种标准配置让研究者在相同条件下对比。3.2 基准任务组合与数据划分从我复现时的理解来看CDFSS基准把现有公开数据集按照“自然图像、遥感图像、医学图像、驾驶场景”几个大类组织起来。每个大类里选择一个或多个数据集作为目标域源域则固定使用包含丰富类别和场景的通用分割数据集。举个我自己跑过的配置源域使用COCO-20i中的训练类别目标域使用LoveDA遥感数据集测试时从目标域每个语义类别随机取若干张作为支持集查询集使用同分布但不相交的图像。这种划分方式对模型不太友好因为LoveDA是高分遥感影像物体尺度、拍摄视角、颜色分布都和自然图像差别极大能更真实地反映模型的跨域能力。下表整理了我在实际复现中常看到的一些域组合可以快速参考源域数据集目标域数据集特点适合验证什么能力PASCAL-5iDeepGlobe自然图像到卫星图像尺度变化和颜色偏移PASCAL-5iISIC 2018自然图像到皮肤镜图像纹理差异和边界模糊COCO-20iLoveDA复杂自然场景到遥感场景类内多样性和域间差异COCO-20iCityscapes夜间子集多种自然场景到车载夜间图像光照变化和运动模糊需要特别强调的是CDFSS基准不是单纯把两个数据集拼起来就行。它有一套标准的episode采样代码包括类别映射、图像尺寸统一、支持集和查询集的划分方式。如果采样逻辑不一致即使使用相同数据集跑出来的mIoU也会差几个点。复现时我用的是论文官方仓库里的自定义Dataset类代码里直接继承了PyTorch的Dataset接口逻辑比较透明。3.3 评估指标与使用建议跨域小样本分割的评估指标主要沿用通用语义分割的标准最常用的是mIoUmean Intersection over Union此外还有F1值、类别平均准确率等。在少样本跨域场景下单看mIoU有时候会产生误导某些类别在目标域中出现频率很低少数几个episode的结果方差很大一次运气好可能就把指标拉上去。论文里给出的解决方案是多次随机采样并报告均值加减标准差。我建议复现时至少跑5次完整的测试流程每次重新采样支持集和查询集然后取平均。这个在代码里实现并不复杂就是外层加个循环记录每次的mIoU最后算均值和方差。如果不这样做对比两个方法时极有可能得出错误结论。另外CDFSS基准还提供了一个“跨域退化率”的参考指标在同一方法上计算在源域测试的mIoU和跨域测试的mIoU之间的相对差距。这个指标可以在一定程度上衡量方法对域偏移的抵抗能力。我在论文的消融实验部分看到用上对齐模块后这个退化率明显降低说明它确实在缓解域偏移方面起了作用。4. 复现教程从环境到指标复现4.1 环境配置与依赖先说复现环境。我用的配置是Ubuntu 20.04、Python 3.8、PyTorch 1.10以上版本、CUDA 11.3、单张RTX 3090。这个项目对显存的要求不算低主要是因为Transformer模块会在特征图上做多头注意力如果输入图像分辨率设到512x512batch又较大显存占用会明显增加。建议创建一个干净的虚拟环境避免依赖冲突。核心依赖如下conda create -n patnet python3.8 conda activate patnet pip install torch1.10.0cu113 torchvision0.11.0cu113 --extra-index-url https://download.pytorch.org/whl/cu113 pip install opencv-python pillow numpy tqdm tensorboard如果使用的是官方仓库还需要安装一些额外的包。仓库里通常有一个requirements.txt直接执行pip install -r requirements.txt即可。这里提醒一点不同机器的CUDA版本不同最好先确认本机驱动支持的CUDA版本再装PyTorch不要盲目用最新版本新版PyTorch有时会改变算子行为导致复现结果和论文对不上。4.2 数据准备与目录组织数据准备是整个复现里最容易出错的一步。CDFSS涉及多个数据集每个数据集的标注格式不一样需要统一预处理。我复现时使用的目录结构如下data/ ├── pascal/ │ ├── JPEGImages/ │ ├── SegmentationClass/ │ └── split/ ├── coco/ │ ├── train2014/ │ ├── val2014/ │ └── annotations/ ├── loveda/ │ ├── train/ │ └── val/ └── deepglobe/ ├── images/ └── masks/每个数据集都需要生成一个类别映射文件把不同数据集中的语义类别对齐到统一的类别编号。比如PASCAL中的“汽车”和LoveDA中的“汽车”可能编号不同需要手动映射。这一步直接影响测试时支持集和查询集的类别对齐一旦映射错误整个测试指标都会乱掉。对于目标域数据如果原图尺寸差异很大建议先做统一的resize或padding。遥感图像通常是几千乘几千的大图直接送进网络不现实。我在实际复现时先把遥感影像裁剪成512x512的patch然后从中采样包含目标类别的patch作为候选支持图。这个处理和原论文的数据准备代码一致裁剪步长设为256保证相邻patch有一定重叠增加有效样本数量。4.3 训练流程与关键超参数训练流程总体上是标准的每个迭代采样一个episode前向传播计算预测和损失反向传播更新参数。具体到PATnet有两个地方值得注意一个是损失函数是多任务组合需要在代码里把分割损失和对齐损失都加上并设置权重另一个是主干网络使用ImageNet预训练权重训练时可以选择冻结部分层也可以全量微调。关键超参数我列了一份参考配置python train.py \ --backbone resnet50 \ --pretrained \ --source_dataset pascal \ --target_dataset loveda \ --n_shot 1 \ --train_episodes 40000 \ --batch_size 4 \ --lr 1e-4 \ --lr_decay step \ --align_weight 0.1 \ --image_size 512这里的align_weight是对齐损失的权重我试过0.05到0.5之间的几个值整体在0.1左右效果比较稳。权重太小对齐模块起不到作用权重太大又会在训练早期干扰分割主干的学习。n_shot设为1表示1-shot如果想跑5-shot就改成5。训练过程中可以开启TensorBoard监控loss曲线。分割loss正常应该在0.6到0.8之间波动如果发现对齐loss持续上升或者剧烈震荡先降低学习率大概率是优化器不稳定不是代码问题。4.4 测试与可视化测试阶段比较重要的是支持集和查询集的采样策略。我建议测试时固定一个随机种子第一次采样后把支持图像和mask保存下来。这样做主要有两个好处一是多个方法之间可以直接对比避免因为随机采样造成的指标波动二是出现问题时可以通过回看支持图确认采样是否正确例如支持图里是否真的包含目标类别mask与图像是否对齐。一个可用的测试命令如下python test.py \ --resume checkpoints/patnet_best.pth \ --test_dataset loveda \ --n_shot 1 \ --eval_episodes 1000 \ --seed 42测试完成后除了在终端打印mIoU还会保存每个类别的IoU。我最开始复现时有一个失误只看平均mIoU没有检查每个类别的IoU。后来发现某个类别在目标域中几乎没有预测成功过但平均mIoU还不错这才意识到单看平均值会有盲区。类别级指标能帮你定位模型在哪些语义上泛化得不好比如模型对“车辆”这类结构清晰的物体效果好对“裸地”这类纹理不明显的类别效果差后续调优方向就清晰了很多。5. 复现中常见问题与调参避坑5.1 典型bug与解决办法复现过程中我踩过不少坑这里挑几个典型的整理成表格遇到类似问题可以先对照排查。现象可能原因解决办法训练loss下降测试mIoU为0support mask没有变成二值索引类别编号错位检查Dataset返回的mask是否只有0和1backbone权重加载报错预训练权重num_classes不匹配删除最后一个全连接层参数再加载显存溢出输入分辨率太大或batch过大把image_size降到384或减少batch_size测试结果波动极大support集采样随机性太大固定随机种子多次采样取平均域对齐loss不收敛权重设置不合适或学习率太高降低对齐损失权重学习率降到5e-5第三个显存溢出问题在Transformer类模型里很常见。我一开始直接把Swin-T作为backbone加上多层Transformer交互模块在3090上也只能跑batch size 2。后来把输入分辨率降到384并把backbone的stage4输出通道压缩问题就解决了。需要注意的是压缩通道要和新增加的对齐模块匹配不是随便改一下就行。5.2 少样本训练中的过拟合问题少样本分割本身就容易过拟合跨域场景更加明显。我观察到一个很典型的现象训练集loss可以降到很低但测试mIoU上不去说明模型只是记住了源域的训练类别没有学会真正的分割能力。一个有效的做法是增加数据增强。PATnet原论文里用了随机缩放、水平翻转、随机颜色扰动这些增强方式看起来简单但在跨域场景下比复杂的数据混合技巧更实用。特别是随机颜色扰动可以直接缩小源域和目标域之间的颜色分布差异对遥感图像迁移帮助很大。还有一个经验是不要盲目增加训练episode数。我试过把训练episode从4万加到8万测试mIoU反而下降了原因是模型在训练后期对源域过拟合。根据我自己的体验可以在训练过程中每隔5000个episode保存一次检查点最后用验证集挑最优检查点而不是简单用最后一个。5.3 域选择的经验之谈训练时加入多个源域并不一定能提升跨域效果。CDFSS基准里经常用COCO作为源域因为COCO类别多、场景多样对提高模型的泛化能力有帮助。但如果你同时混入PASCAL和COCO由于两个数据集标注风格和类别定义不同反而会让模型在特征空间里产生冲突。我自己的建议是在复现原论文时先严格按照论文的配置跑一遍把基线指标复现出来然后再做你自己的实验比如尝试不同的源域组合。直接一上来就改配置出了问题很难判断是方法问题还是配置问题。另一个实用的技巧是“域代理”在训练前用一个小型特征提取器分别编码源域和目标域的图像观察特征分布之间的距离。如果两个分布几乎完全不重叠说明域差异太大即使PATnet也未必能救回来如果分布有部分重叠域对齐模块的收益最明显。这个预处理不需要太复杂的工具用t-SNE可视化特征分布就够了。6. 个人体会与可扩展的思路6.1 我复现过程中的最大感受做完整轮复现后我感受最深的是域对齐做进Transformer交互这个设计的巧妙之处。常规的域适应方法通常要求训练时能同时看到源域和目标域数据但PATnet把对齐过程化解为支持原型和查询特征的注意力交互让模型在推理时天然具备跨域能力。这种思路在工程上的意义在于训练和推理是同一个模型不需要在目标域上二次适配部署成本低了很多。另外CDFSS基准的规范价值也值得肯定。之前我在对比不同方法时最大的烦恼是各方法实验设置不统一需要花大量时间重新实现别人的代码。有了统一基准后至少大家是在同一个起跑线上比较论文里的指标也更有参考意义。6.2 可以从这个工作延伸的方向从PATnet往后看我觉得有几个方向值得继续尝试。一个是和大规模预训练模型结合。现在SAM这类视觉基础模型在很多分割任务上表现出强大的泛化能力如果能用SAM生成伪标签或提供更稳定的特征表示再结合PATnet的对齐机制也许能进一步减少对目标域标注的依赖。另一个是动态原型。PATnet中支持原型在每次推理时是固定的但在目标域类别外观差异极大时单一原型仍然可能不够。后续可以考虑根据查询特征动态调整原型比如在测试阶段用一次前向传播生成粗略分割结果再用粗略结果重建更准确的原型形成两阶段迭代。最后是扩展到点云或视频域。跨域问题不只存在于图像中点云语义分割在不同雷达配置之间同样存在严重的域偏移。PATnet中的原型对齐思路能不能迁移到3D分支值得做一做实验。这个工作的完整复现过程代码、数据、参数踩坑我都拆开讲了一遍核心就是一句话跨域小样本分割的难点不在网络结构有多深而在于怎么让模型在极少量目标域样本面前依然保持稳定的语义理解能力。把这个思路吃透后续不管换数据集还是换骨干网络都不会跑偏。