做语义分割的同行应该都有同感RGB-D任务里最难的不是分割网络本身而是怎么让RGB和Depth两个模态真正“配合”起来。纯RGB分割你已经玩得很熟了可一旦加入深度图很多人会发现精度反而可能下降——因为深度图噪声大、语义信息弱如果跟RGB做对称融合它更像是在往特征里掺沙子。PrimKD这个工作核心思路就是打破这种对称提出“主模态引导”的多模态融合思路让RGB分支充当teacher的角色再用知识蒸馏去约束Depth分支的学习方向最终在RGB-D语义分割上拿到更稳定、更一致的精度提升。这篇文章我想从原理层面拆开讲清楚这套设计的内在逻辑再把我自己在复现和调参过程中验证过的实操细节、超参设置、以及踩过的坑一并整理出来。不管你是刚接触多模态分割还是已经做过一段时间RGB-D方向这篇文章应该都能提供一些可以直接落地的参考。1. 项目背景RGB-D语义分割为什么需要多模态融合1.1 任务定义不只有相机还有距离RGB-D语义分割的输入是两张图一张常规的彩色图像记录的是颜色和纹理一张深度图像记录的是每个像素到相机的物理距离。模型要输出的是一张像素级的语义标签图每个像素属于哪个类别——桌子、椅子、墙面、人、杯子等等。跟纯RGB分割相比RGB-D最大的优势在于深度模态天然携带了三维几何信息可以帮助网络更好地区分那些纹理相似但空间位置和形状不同的物体。现实应用里这类需求非常具体。机器人抓取时机械臂需要知道“这个杯子的边缘在哪里”而不是只知道“这里有一个杯子的像素集合”AR导航时系统需要把虚拟信息稳定叠在真实物体表面如果分割边界忽忽悠悠就完全不可用在室内SLAM和辅助驾驶场景中深度信息还能在暗光、强反光这类RGB几乎罢工的条件下提供至少轮廓级别的线索。但深度图并不像一个“友好”的输入信号。传感器在黑色物体、透明物体、远处区域上经常直接给到无效值深度噪声会让物体边缘看起来像是被啃过一样。更麻烦的是深度图本身不具备语义区分能力彩色上完全一样的物体可能在不同距离出现深度值一样但完全不等价。所以RGB-D分割的难点从来不是“把两个输入拼起来”而是怎样让两种完全不同性质的信息互相配合而不互相污染。1.2 多模态融合的传统做法与困境早期做法非常直接把RGB和Depth拼成四通道图像喂给网络。这种方法属于early fusion处理起来方便但效果普遍一般。原因不难理解RGB和Depth的统计特性差异太大了一个接近自然图像分布另一个是距离场分布在输入层面强行叠在一起卷积核很难同时提取出两类特征的有效部分反而容易让网络把深度噪声也当成有效纹理学进去。后来主流做法演变成双分支结构RGB分支和Depth分支各自用一套卷积提取特征再在中间层或最后层做融合。这类方案里dense fusion是效果最稳定的一档——在多个尺度上都做特征融合而不是只靠最后一层。但即便到了这一步依然有一个基础假设没被质疑过RGB和Depth在融合中是平等、对称的。事实并不是这样。RGB在绝大多数场景中承担了绝大部分语义判断类别边界、物体属性、场景类型这些信息基本都由颜色纹理决定Depth则主要在物体边界、空间形状、尺度关系上提供补充。对称融合的问题在于网络必须自己摸索出“哪些深度特征该信、哪些不该信”的规则而深度图的噪声分布是高度不均匀的网络很容易在某些样本上学偏把深度图的伪影当作稳定的语义线索。2. 核心思路拆解主模态引导与蒸馏的配合逻辑2.1 RGB为什么天然是“主模态”多模态融合里有一个经常被忽略的事实不同模态的信息密度和可靠性天然不一样。RGB和Depth不是两个对等的通道RGB是“主”Depth是“辅”。这不是拍脑袋而是任务本身决定的。拿室内场景举例。一张红色沙发一张红色木桌在RGB里可以用纹理和形制轻松分开但只给深度图的话这两个物体的距离、形状可能非常接近单纯靠深度根本无法区分。反过来Depth提供的物体边界和空间轮廓虽然语义上“说什么类别”不给力但在“哪里换到下一个物体”这件事上准确率往往比RGB还高。也就是说RGB负责“是什么”Depth负责“在哪切”这两个职责完全不对等。这就像两个人合作辨认一个物体一个人视力好能看清颜色、花纹、文字另一个人负责摸轮廓和距离。视力好的人必须做决策主导另一个人提供几何线索来辅助判断。如果让两个人完全平等地投票视力差的那个人会把很多错误判断带进来。PrimKD的核心想法就是把这种不对称性显式建模RGB分支是primary modality融合时占据主导地位Depth分支是secondary modality用来补充几何边界信息。2.2 Knowledge Distillation在主框架中扮演什么角色这里有个非常关键的设计问题既然指定了RGB当主模态那Depth分支怎么办直接砍掉深度信息就回到了单模态不砍又怕它干扰主导特征。PrimKD的答案是用知识蒸馏让Depth分支向RGB分支的特征空间对齐。知识蒸馏通常被用在模型压缩上用一个大的教师网络教小的学生网络。但这里的用法完全不同。PrimKD里的KD不是为了压缩模型而是为了解决多模态特征表示不一致的问题。RGB分支学到的特征空间偏向语义类别Depth分支学到的特征空间偏向几何结构两个分布差异太大融合时简单相加或concat只会互相抵消。通过蒸馏loss把RGB分支在某个中间层的特征作为“软目标”监督Depth分支对应层的特征去学习。不是说让Depth丢失几何信息而是让它在保留空间细节的前提下把特征分布向语义方向拉近。这样融合模块拿到的两个分支特征处于同一个语义特征空间中一个提供语义主体一个提供几何修正融合效果自然比两个各自为政的分支要好得多。这个思路在实操中还有一个额外收益它天然具备一定的抗噪能力。深度图的噪声区域产生的特征如果和RGB分支的语义特征差异过大蒸馏loss会把这些方向上的响应压下去相当于给Depth分支装了一个“去伪存真”的约束。2.3 整体信息流与训练目标整个框架的数据流可以这样描述RGB图像输入主分支Depth图像输入辅助分支两个分支在多个尺度提取特征辅助分支的特征在蒸馏loss的监督下向主分支对齐逐层融合后的特征送入解码器预测分割图。训练目标由三部分组成L L_seg α · L_distill β · L_aux其中L_seg是分割主loss通常用交叉熵或其变体L_distill是主模态对辅助模态的蒸馏约束L_aux是辅助分支自己的辅助loss可选能帮助Depth分支更快掌握几何结构的表征。α和β是两个权重超参需要根据训练情况动态调整这个我在后面实操章节会具体展开说明。从整体看PrimKD并不是一个特别重的改造它没有发明全新的复杂融合模块而是通过“谁是主、谁是辅”这个切入点把蒸馏和融合重新组织了一遍。这种思路迁移性很强理论上任何双分支融合结构都能套进这个框架里。3. 实操实现数据准备、网络搭建与损失配置3.1 数据集与预处理细节我在复现时主要使用了室内RGB-D分割的常用公开数据集比如NYUDv2和SUN RGB-D这类。两者的标注类别、场景范围、深度图质量差异都不小正好可以检验方法是不是只在单一benchmark上有效。预处理是最容易被忽略的环节。原生的深度图一般是16位整型存的是毫米或厘米单位的距离值。直接归一化到[0,1]会让网络学得很痛苦因为深度数据的分布极度偏向近距离区域远距离的值出现频率很低。常规做法是先做裁剪把超出传感器量程的深度值截断比如保距在0到8米或0到10米之间的像素再归一化。裁剪后还要过一个log变换因为场景中深度值的物理意义是空间距离人眼感知的距离和线性深度不成比例取log能让网络更容易学到距离带来的尺度差异。对深度图做编码时HHA编码是我权衡后最推荐的选择。HHA把深度图拆成水平视差、高度、垂直视差三个通道相当于从深度图中显式还原了部分几何结构线索网络不用完全自己强行从单通道距离图里学习。代价是预处理会多花一些时间但换来的是训练收敛速度和精度的双重收益。数据增强环节要格外注意RGB和Depth必须使用完全相同的几何变换随机翻转、随机缩放、随机裁剪都要用同一套参数否则两个模态的空间对应关系被破坏训练直接白费。我自己踩过一次坑只对RGB做了随机颜色抖动Depth图没做虽然不会破坏空间对齐但会导致两个分支看到的输入统计特性始终不同蒸馏loss在前几个epoch非常不稳定。后来把颜色抖动也去掉问题立刻缓解。3.2 骨干网络与融合模块的实现骨干网络部分我选择了可插拔设计主干可以用常见的ResNet系列也可以换成基于Transformer的分层编码器。PrimKD本身不绑定特定backbone这对实际落地来说非常友好。如果你有现成的预训练权重直接用就行如果没有从ImageNet预训练开始也能收敛。如果使用ResNet作为骨干RGB和Depth各加载一套backbone权重。中间层特征引出后在多个尺度上和辅助分支做融合。这里给出一个我在复现时使用的简化融合模块核心是主模态引导的门控结构import torch import torch.nn as nn class PrimaryGuidedFusion(nn.Module): def __init__(self, in_channels, hidden_channels256): super().__init__() self.rgb_proj nn.Conv2d(in_channels, hidden_channels, 1) self.depth_proj nn.Conv2d(in_channels, hidden_channels, 1) self.fuse nn.Sequential( nn.Conv2d(hidden_channels * 2, hidden_channels, 3, padding1), nn.BatchNorm2d(hidden_channels), nn.ReLU(inplaceTrue), ) self.gate nn.Sequential( nn.Conv2d(hidden_channels * 2, hidden_channels, 3, padding1), nn.Sigmoid(), ) def forward(self, rgb_feat, depth_feat): rgb self.rgb_proj(rgb_feat) depth self.depth_proj(depth_feat) cat_feat torch.cat([rgb, depth], dim1) gate self.gate(cat_feat) fused self.fuse(cat_feat) # 主模态引导融合结果通过门控补充到RGB主特征上 return fused * gate rgb这个模块的思路可以这样理解fuse分支从RGB和Depth拼接特征中生成一个融合候选gate分支负责判断每个空间位置上该信任多少深度信息最终把经过门控的融合候选累加到RGB主特征上。因为始终保留了rgb特征本身主模态的作用就被显式保证。实际操作中这个模块会插入到backbone的多个stage之后实现多尺度的渐进融合。值得提醒的是这里为了可读性展示的是最简版本照搬到真正训练时会有几个问题。一是单层融合感受野不够建议在更深stage使用膨胀卷积或者更大kernel二是BatchNorm在小batch size下不稳定如果显存只够跑很小的batch建议换成GroupNorm或同步BN。3.3 蒸馏损失的具体形态蒸馏loss有两种常用形态feature-level和logits-level。PrimKD这类方法更适合用feature-level因为目标是让辅助分支的中间特征向主分支对齐而不只是最终的类别概率一致。特征对齐最朴素的做法是直接算两个特征图的MSE但这里有一个问题两个分支的backbone如果一样特征图空间尺寸相同但通道语义不一定对齐。我建议在Depth分支特征后接一个1x1的投影卷积把通道数映射到和RGB分支一致再做匹配。投影层本身参与训练让蒸馏过程有足够的变换自由度。logits-level蒸馏也有价值它施加在分类输出上能够让Depth分支独立预测的语义结果也保持和RGB分支一致。这时候温度T是一个关键超参。常规设置T4.0蒸馏loss变成如下形式import torch.nn.functional as F def distill_loss(teacher_logits, student_logits, temperature4.0): t temperature teacher_prob F.log_softmax(teacher_logits / t, dim1) student_prob F.log_softmax(student_logits / t, dim1) loss F.kl_div(student_prob, teacher_prob, reductionbatchmean) return loss * (t ** 2)乘上T²是标准KD的常见处理目的是让梯度尺度不会因为温度缩放过大或过小。为什么要用log_softmax而不是直接对logits算MSE因为分类输出经过softmax之后真正有意义的是类别之间的相对关系KL散度能捕捉这种相对分布差异而MSE对整体幅度的变化太过敏感。实操中还有个容易忽略的点temperature并不是越大越好。T过大Softmax输出趋于均匀分布蒸馏目标变成“所有类别概率都差不多”辅助分支学不到精细的类别边界T过小蒸馏又退化成hard label和直接用one-hot没区别。我在实验中的经验是先从T4开始训练几个epoch观察蒸馏loss曲线如果下降太慢再往6调如果辅助分支很快过拟合到教师输出、但融合后精度没提升就往回降到2。4. 训练调参与性能评估要点4.1 核心训练配置速查我这里整理了一份自己在复现过程中实际验证过的训练配置不一定对所有数据集都最优但作为初始基线非常可靠。配置项推荐值备注优化器AdamW比SGD收敛稳定配合权重衰减效果好初始学习率6e-5batch 16学习率随batch size线性缩放学习率策略Poly幂次0.9分割任务比cosine更常用权重衰减0.01AdamW默认推荐训练轮数300~400多模态训练收敛更慢急不得蒸馏权重α前50轮0.3之后衰减到0.1前期靠蒸馏快速对齐后期不能让蒸馏压制辅助分支特性辅助loss权重β0.1帮助Depth分支稳定早期训练混合精度FP16显存和速度都能优化约40%这里很关键的一个设置是蒸馏权重α的退火。如果全程固定一个较大的αDepth分支会被训练成RGB分支的劣化复制品其独有的几何信息会被蒸馏约束压制掉。前期α大是为了让辅助分支快速学会哪些特征是语义有效的后期α减小是为了释放自由度让深度分支去利用自身模态独有的几何线索。这种先对齐再放开的两阶段策略在我所有对比实验里都稳定优于固定权重。4.2 评估指标与对照实验设计评估指标主要看mIoUmean Intersection over Union像素准确率pixAcc作为参考即可。mIoU对边界质量和类别不平衡更敏感能更好地区分不同融合方法的真实差距。做方法对比实验时我强烈建议至少准备三组基线纯RGB分支不输入深度这是最重要的信息上限参考纯Depth分支几乎没人用它做最终分割但能帮你了解深度模态单独的性能下限双分支concat融合不加主模态引导和蒸馏这是最直接的同类比较。我完成的一组对比实验结果显示在相同backbone设置下双分支concat融合相对纯RGB在部分类别上有提升但在细纹理类别上出现过下降加入主模态引导后下降被明显抑制再加入蒸馏loss后整体mIoU进一步提升尤其在物体边界区域的分割质量改善非常明显。需要说明的是不同数据集的收益幅度不一致深度图质量好的数据集上提升更显著深度图本身很烂的场景中提升幅度会收窄但基本不会出现比双分支基线更差的情况。还有一个值得观察的现象蒸馏loss收敛后把辅助分支单独拆出来跑分割性能虽然比不上RGB分支但比从头训练的Depth-only模型要高明得多。这从侧面说明蒸馏并没有消灭Depth的几何能力反而是在保留几何线索的同时让它的语义表达能力更强了。4.3 多模态训练的小技巧多模态训练的显存开销几乎是单模态的两倍这是最现实的工程问题。我的做法是把backbone拆开处理前几个stage的特征已经足够低级不需要两套完整backbone都跑只在较深stage分叉出Depth辅助分支这样能省一大部分显存。如果你坚持在浅层就分离出双分支建议至少冻结RGB主干的前两个stage让模型优先学习融合和语义解码部分。训练初期还有一个常见的坑两个分支如果学习率一致Depth分支往往更新得比RGB分支慢因为深度特征上的梯度普遍更小。我倾向于把Depth分支的初始学习率调成RGB分支的0.8倍再配合蒸馏loss。这样两者更新步调不至于差太远融合模块收到的两个输入的质量也相对均衡。5. 常见问题与排查技巧实录5.1 深度图缺失和噪声怎么处理深度图在黑色物体、透明材质、高光表面以及超出传感器量程的区域经常会直接变成无效值。很多新手拿到数据后第一反应是把缺失像素置为0。这个操作虽然可以让训练跑通但会引入一个隐患网络会学到“深度值为0的位置就是物体边界”而这个规律在真实场景中并不成立因为远距离区域的深度也可能被截断到0附近。我的处理方式是在输入层做两步。第一步缺失像素统一重设为一个“远离量程”的常量比如量程是8米缺失就设为8.5然后归一化时不会和其他正常深度混在一起第二步对缺失区域生成一个valid mask输入给Depth分支让网络明确知道哪些位置的深度不可信。实验下来这个组合比直接把缺失区域抠掉或者补成0都要稳定。训练时还可以做一种随机mask增强随机丢弃某个区域内的深度值强迫模型学习在部分深度缺失时仍然能保持分割质量。实测下来这种增强能显著提升模型在真实传感器上的鲁棒性因为开源数据集的深度质量普遍比自采数据好不做这个增强的话模型很容易出现过拟合。5.2 蒸馏信号不稳定怎么办特征蒸馏loss的数值曲线如果出现大幅震荡往往不是方法问题而是设备问题。最常见的原因是特征匹配层的数值量级不对齐一个分支的输出值域是[0,1]级别另一个是[-10,10]级别直接算MSE会让梯度被绝对值大的那个分支带偏。在投影层后加一个LayerNorm或BatchNorm强制把特征归一化到相同尺度震荡问题基本能解决大半。还有一个隐蔽问题就是两个分支的训练进度不一致。如果RGB分支已经收敛到不错的状态而Depth分支还在早期挣扎蒸馏loss会给Depth分支非常大的梯度甚至有把已经训练好的RGB特征也给带坏的情况。这就是为什么我前面建议在RGB主干相对稳定之前可以先冻结主干只训练融合模块和辅助分支等辅助分支有一定表达能力了再解冻主干联合训练。如果logits蒸馏的KL散度在训练中期开始变成一条水平线说明辅助分支已经学不到更多来自教师的信息了。这时候可以去检查温度系数是不是太小概率分布太尖锐梯度几乎只在真实类别那一维上有值其他维度全是0学不到软标签里的类别关系信息。5.3 工程效率与显存优化双backbone结构最直接的问题就是显存开销接近翻倍。我实际测试下来输入分辨率如果跑到512x512batch size设8单卡16G是基本门槛。优化手段有几种深度分支选用更轻量的结构不一定要和RGB共享同一个backboneDepth特征不需要那么高的语义抽象能力轻量网络足够在stage之间插入融合点而不是每个stage都融合融合点越多梯度路径越复杂显存占用也越多3个融合点是一个比较均衡的选择使用梯度累积等效扩大batch size能有效稳定BN层统计量多模态训练对batch size的敏感度比单模态更高batch太小BN跑偏的概率大幅增加。另一个工程细节是Depth和RGB的输入分辨率对齐问题。有些公开数据集的深度图和RGB图尺寸不一致DataLoader里必须做好resize和crop的统一。千万别用两套不同的resize参数否则后面所有工作都白做。写在最后我最早复现这套主模态引导方案时其实是持怀疑态度的觉得无非是给深度分支加了点约束能有多大区别。实际跑下来才知道多模态融合真正的问题往往不是出在融合模块本身而是两个分支的特征空间能不能顺畅对话。PrimKD不搞复杂的注意力花活用一个蒸馏约束把两个分支的表示对齐了后面所有融合模块都变得好训练了。那个“先对齐再放开”的训练策略是我在整个复现过程中收获最大的经验。反过来想这套“主模态蒸馏”的思路其实可以迁移到很多模态组合上RGB加LiDAR点云、RGB加事件相机、语言加视觉都存在类似的信息不对称和特征空间不匹配问题。只要你能找到那个信息密度最高的主模态再用蒸馏去引导辅助模态对齐大概率都能得到一组更好的融合效果。说到底多模态融合不应该是把所有输入一视同仁地塞进同一个网络而是搞清楚每个模态擅长什么再让它们各司其职。
做语义分割的同行应该都有同感RGB-D任务里最难的不是分割网络本身而是怎么让RGB和Depth两个模态真正“配合”起来。纯RGB分割你已经玩得很熟了可一旦加入深度图很多人会发现精度反而可能下降——因为深度图噪声大、语义信息弱如果跟RGB做对称融合它更像是在往特征里掺沙子。PrimKD这个工作核心思路就是打破这种对称提出“主模态引导”的多模态融合思路让RGB分支充当teacher的角色再用知识蒸馏去约束Depth分支的学习方向最终在RGB-D语义分割上拿到更稳定、更一致的精度提升。这篇文章我想从原理层面拆开讲清楚这套设计的内在逻辑再把我自己在复现和调参过程中验证过的实操细节、超参设置、以及踩过的坑一并整理出来。不管你是刚接触多模态分割还是已经做过一段时间RGB-D方向这篇文章应该都能提供一些可以直接落地的参考。1. 项目背景RGB-D语义分割为什么需要多模态融合1.1 任务定义不只有相机还有距离RGB-D语义分割的输入是两张图一张常规的彩色图像记录的是颜色和纹理一张深度图像记录的是每个像素到相机的物理距离。模型要输出的是一张像素级的语义标签图每个像素属于哪个类别——桌子、椅子、墙面、人、杯子等等。跟纯RGB分割相比RGB-D最大的优势在于深度模态天然携带了三维几何信息可以帮助网络更好地区分那些纹理相似但空间位置和形状不同的物体。现实应用里这类需求非常具体。机器人抓取时机械臂需要知道“这个杯子的边缘在哪里”而不是只知道“这里有一个杯子的像素集合”AR导航时系统需要把虚拟信息稳定叠在真实物体表面如果分割边界忽忽悠悠就完全不可用在室内SLAM和辅助驾驶场景中深度信息还能在暗光、强反光这类RGB几乎罢工的条件下提供至少轮廓级别的线索。但深度图并不像一个“友好”的输入信号。传感器在黑色物体、透明物体、远处区域上经常直接给到无效值深度噪声会让物体边缘看起来像是被啃过一样。更麻烦的是深度图本身不具备语义区分能力彩色上完全一样的物体可能在不同距离出现深度值一样但完全不等价。所以RGB-D分割的难点从来不是“把两个输入拼起来”而是怎样让两种完全不同性质的信息互相配合而不互相污染。1.2 多模态融合的传统做法与困境早期做法非常直接把RGB和Depth拼成四通道图像喂给网络。这种方法属于early fusion处理起来方便但效果普遍一般。原因不难理解RGB和Depth的统计特性差异太大了一个接近自然图像分布另一个是距离场分布在输入层面强行叠在一起卷积核很难同时提取出两类特征的有效部分反而容易让网络把深度噪声也当成有效纹理学进去。后来主流做法演变成双分支结构RGB分支和Depth分支各自用一套卷积提取特征再在中间层或最后层做融合。这类方案里dense fusion是效果最稳定的一档——在多个尺度上都做特征融合而不是只靠最后一层。但即便到了这一步依然有一个基础假设没被质疑过RGB和Depth在融合中是平等、对称的。事实并不是这样。RGB在绝大多数场景中承担了绝大部分语义判断类别边界、物体属性、场景类型这些信息基本都由颜色纹理决定Depth则主要在物体边界、空间形状、尺度关系上提供补充。对称融合的问题在于网络必须自己摸索出“哪些深度特征该信、哪些不该信”的规则而深度图的噪声分布是高度不均匀的网络很容易在某些样本上学偏把深度图的伪影当作稳定的语义线索。2. 核心思路拆解主模态引导与蒸馏的配合逻辑2.1 RGB为什么天然是“主模态”多模态融合里有一个经常被忽略的事实不同模态的信息密度和可靠性天然不一样。RGB和Depth不是两个对等的通道RGB是“主”Depth是“辅”。这不是拍脑袋而是任务本身决定的。拿室内场景举例。一张红色沙发一张红色木桌在RGB里可以用纹理和形制轻松分开但只给深度图的话这两个物体的距离、形状可能非常接近单纯靠深度根本无法区分。反过来Depth提供的物体边界和空间轮廓虽然语义上“说什么类别”不给力但在“哪里换到下一个物体”这件事上准确率往往比RGB还高。也就是说RGB负责“是什么”Depth负责“在哪切”这两个职责完全不对等。这就像两个人合作辨认一个物体一个人视力好能看清颜色、花纹、文字另一个人负责摸轮廓和距离。视力好的人必须做决策主导另一个人提供几何线索来辅助判断。如果让两个人完全平等地投票视力差的那个人会把很多错误判断带进来。PrimKD的核心想法就是把这种不对称性显式建模RGB分支是primary modality融合时占据主导地位Depth分支是secondary modality用来补充几何边界信息。2.2 Knowledge Distillation在主框架中扮演什么角色这里有个非常关键的设计问题既然指定了RGB当主模态那Depth分支怎么办直接砍掉深度信息就回到了单模态不砍又怕它干扰主导特征。PrimKD的答案是用知识蒸馏让Depth分支向RGB分支的特征空间对齐。知识蒸馏通常被用在模型压缩上用一个大的教师网络教小的学生网络。但这里的用法完全不同。PrimKD里的KD不是为了压缩模型而是为了解决多模态特征表示不一致的问题。RGB分支学到的特征空间偏向语义类别Depth分支学到的特征空间偏向几何结构两个分布差异太大融合时简单相加或concat只会互相抵消。通过蒸馏loss把RGB分支在某个中间层的特征作为“软目标”监督Depth分支对应层的特征去学习。不是说让Depth丢失几何信息而是让它在保留空间细节的前提下把特征分布向语义方向拉近。这样融合模块拿到的两个分支特征处于同一个语义特征空间中一个提供语义主体一个提供几何修正融合效果自然比两个各自为政的分支要好得多。这个思路在实操中还有一个额外收益它天然具备一定的抗噪能力。深度图的噪声区域产生的特征如果和RGB分支的语义特征差异过大蒸馏loss会把这些方向上的响应压下去相当于给Depth分支装了一个“去伪存真”的约束。2.3 整体信息流与训练目标整个框架的数据流可以这样描述RGB图像输入主分支Depth图像输入辅助分支两个分支在多个尺度提取特征辅助分支的特征在蒸馏loss的监督下向主分支对齐逐层融合后的特征送入解码器预测分割图。训练目标由三部分组成L L_seg α · L_distill β · L_aux其中L_seg是分割主loss通常用交叉熵或其变体L_distill是主模态对辅助模态的蒸馏约束L_aux是辅助分支自己的辅助loss可选能帮助Depth分支更快掌握几何结构的表征。α和β是两个权重超参需要根据训练情况动态调整这个我在后面实操章节会具体展开说明。从整体看PrimKD并不是一个特别重的改造它没有发明全新的复杂融合模块而是通过“谁是主、谁是辅”这个切入点把蒸馏和融合重新组织了一遍。这种思路迁移性很强理论上任何双分支融合结构都能套进这个框架里。3. 实操实现数据准备、网络搭建与损失配置3.1 数据集与预处理细节我在复现时主要使用了室内RGB-D分割的常用公开数据集比如NYUDv2和SUN RGB-D这类。两者的标注类别、场景范围、深度图质量差异都不小正好可以检验方法是不是只在单一benchmark上有效。预处理是最容易被忽略的环节。原生的深度图一般是16位整型存的是毫米或厘米单位的距离值。直接归一化到[0,1]会让网络学得很痛苦因为深度数据的分布极度偏向近距离区域远距离的值出现频率很低。常规做法是先做裁剪把超出传感器量程的深度值截断比如保距在0到8米或0到10米之间的像素再归一化。裁剪后还要过一个log变换因为场景中深度值的物理意义是空间距离人眼感知的距离和线性深度不成比例取log能让网络更容易学到距离带来的尺度差异。对深度图做编码时HHA编码是我权衡后最推荐的选择。HHA把深度图拆成水平视差、高度、垂直视差三个通道相当于从深度图中显式还原了部分几何结构线索网络不用完全自己强行从单通道距离图里学习。代价是预处理会多花一些时间但换来的是训练收敛速度和精度的双重收益。数据增强环节要格外注意RGB和Depth必须使用完全相同的几何变换随机翻转、随机缩放、随机裁剪都要用同一套参数否则两个模态的空间对应关系被破坏训练直接白费。我自己踩过一次坑只对RGB做了随机颜色抖动Depth图没做虽然不会破坏空间对齐但会导致两个分支看到的输入统计特性始终不同蒸馏loss在前几个epoch非常不稳定。后来把颜色抖动也去掉问题立刻缓解。3.2 骨干网络与融合模块的实现骨干网络部分我选择了可插拔设计主干可以用常见的ResNet系列也可以换成基于Transformer的分层编码器。PrimKD本身不绑定特定backbone这对实际落地来说非常友好。如果你有现成的预训练权重直接用就行如果没有从ImageNet预训练开始也能收敛。如果使用ResNet作为骨干RGB和Depth各加载一套backbone权重。中间层特征引出后在多个尺度上和辅助分支做融合。这里给出一个我在复现时使用的简化融合模块核心是主模态引导的门控结构import torch import torch.nn as nn class PrimaryGuidedFusion(nn.Module): def __init__(self, in_channels, hidden_channels256): super().__init__() self.rgb_proj nn.Conv2d(in_channels, hidden_channels, 1) self.depth_proj nn.Conv2d(in_channels, hidden_channels, 1) self.fuse nn.Sequential( nn.Conv2d(hidden_channels * 2, hidden_channels, 3, padding1), nn.BatchNorm2d(hidden_channels), nn.ReLU(inplaceTrue), ) self.gate nn.Sequential( nn.Conv2d(hidden_channels * 2, hidden_channels, 3, padding1), nn.Sigmoid(), ) def forward(self, rgb_feat, depth_feat): rgb self.rgb_proj(rgb_feat) depth self.depth_proj(depth_feat) cat_feat torch.cat([rgb, depth], dim1) gate self.gate(cat_feat) fused self.fuse(cat_feat) # 主模态引导融合结果通过门控补充到RGB主特征上 return fused * gate rgb这个模块的思路可以这样理解fuse分支从RGB和Depth拼接特征中生成一个融合候选gate分支负责判断每个空间位置上该信任多少深度信息最终把经过门控的融合候选累加到RGB主特征上。因为始终保留了rgb特征本身主模态的作用就被显式保证。实际操作中这个模块会插入到backbone的多个stage之后实现多尺度的渐进融合。值得提醒的是这里为了可读性展示的是最简版本照搬到真正训练时会有几个问题。一是单层融合感受野不够建议在更深stage使用膨胀卷积或者更大kernel二是BatchNorm在小batch size下不稳定如果显存只够跑很小的batch建议换成GroupNorm或同步BN。3.3 蒸馏损失的具体形态蒸馏loss有两种常用形态feature-level和logits-level。PrimKD这类方法更适合用feature-level因为目标是让辅助分支的中间特征向主分支对齐而不只是最终的类别概率一致。特征对齐最朴素的做法是直接算两个特征图的MSE但这里有一个问题两个分支的backbone如果一样特征图空间尺寸相同但通道语义不一定对齐。我建议在Depth分支特征后接一个1x1的投影卷积把通道数映射到和RGB分支一致再做匹配。投影层本身参与训练让蒸馏过程有足够的变换自由度。logits-level蒸馏也有价值它施加在分类输出上能够让Depth分支独立预测的语义结果也保持和RGB分支一致。这时候温度T是一个关键超参。常规设置T4.0蒸馏loss变成如下形式import torch.nn.functional as F def distill_loss(teacher_logits, student_logits, temperature4.0): t temperature teacher_prob F.log_softmax(teacher_logits / t, dim1) student_prob F.log_softmax(student_logits / t, dim1) loss F.kl_div(student_prob, teacher_prob, reductionbatchmean) return loss * (t ** 2)乘上T²是标准KD的常见处理目的是让梯度尺度不会因为温度缩放过大或过小。为什么要用log_softmax而不是直接对logits算MSE因为分类输出经过softmax之后真正有意义的是类别之间的相对关系KL散度能捕捉这种相对分布差异而MSE对整体幅度的变化太过敏感。实操中还有个容易忽略的点temperature并不是越大越好。T过大Softmax输出趋于均匀分布蒸馏目标变成“所有类别概率都差不多”辅助分支学不到精细的类别边界T过小蒸馏又退化成hard label和直接用one-hot没区别。我在实验中的经验是先从T4开始训练几个epoch观察蒸馏loss曲线如果下降太慢再往6调如果辅助分支很快过拟合到教师输出、但融合后精度没提升就往回降到2。4. 训练调参与性能评估要点4.1 核心训练配置速查我这里整理了一份自己在复现过程中实际验证过的训练配置不一定对所有数据集都最优但作为初始基线非常可靠。配置项推荐值备注优化器AdamW比SGD收敛稳定配合权重衰减效果好初始学习率6e-5batch 16学习率随batch size线性缩放学习率策略Poly幂次0.9分割任务比cosine更常用权重衰减0.01AdamW默认推荐训练轮数300~400多模态训练收敛更慢急不得蒸馏权重α前50轮0.3之后衰减到0.1前期靠蒸馏快速对齐后期不能让蒸馏压制辅助分支特性辅助loss权重β0.1帮助Depth分支稳定早期训练混合精度FP16显存和速度都能优化约40%这里很关键的一个设置是蒸馏权重α的退火。如果全程固定一个较大的αDepth分支会被训练成RGB分支的劣化复制品其独有的几何信息会被蒸馏约束压制掉。前期α大是为了让辅助分支快速学会哪些特征是语义有效的后期α减小是为了释放自由度让深度分支去利用自身模态独有的几何线索。这种先对齐再放开的两阶段策略在我所有对比实验里都稳定优于固定权重。4.2 评估指标与对照实验设计评估指标主要看mIoUmean Intersection over Union像素准确率pixAcc作为参考即可。mIoU对边界质量和类别不平衡更敏感能更好地区分不同融合方法的真实差距。做方法对比实验时我强烈建议至少准备三组基线纯RGB分支不输入深度这是最重要的信息上限参考纯Depth分支几乎没人用它做最终分割但能帮你了解深度模态单独的性能下限双分支concat融合不加主模态引导和蒸馏这是最直接的同类比较。我完成的一组对比实验结果显示在相同backbone设置下双分支concat融合相对纯RGB在部分类别上有提升但在细纹理类别上出现过下降加入主模态引导后下降被明显抑制再加入蒸馏loss后整体mIoU进一步提升尤其在物体边界区域的分割质量改善非常明显。需要说明的是不同数据集的收益幅度不一致深度图质量好的数据集上提升更显著深度图本身很烂的场景中提升幅度会收窄但基本不会出现比双分支基线更差的情况。还有一个值得观察的现象蒸馏loss收敛后把辅助分支单独拆出来跑分割性能虽然比不上RGB分支但比从头训练的Depth-only模型要高明得多。这从侧面说明蒸馏并没有消灭Depth的几何能力反而是在保留几何线索的同时让它的语义表达能力更强了。4.3 多模态训练的小技巧多模态训练的显存开销几乎是单模态的两倍这是最现实的工程问题。我的做法是把backbone拆开处理前几个stage的特征已经足够低级不需要两套完整backbone都跑只在较深stage分叉出Depth辅助分支这样能省一大部分显存。如果你坚持在浅层就分离出双分支建议至少冻结RGB主干的前两个stage让模型优先学习融合和语义解码部分。训练初期还有一个常见的坑两个分支如果学习率一致Depth分支往往更新得比RGB分支慢因为深度特征上的梯度普遍更小。我倾向于把Depth分支的初始学习率调成RGB分支的0.8倍再配合蒸馏loss。这样两者更新步调不至于差太远融合模块收到的两个输入的质量也相对均衡。5. 常见问题与排查技巧实录5.1 深度图缺失和噪声怎么处理深度图在黑色物体、透明材质、高光表面以及超出传感器量程的区域经常会直接变成无效值。很多新手拿到数据后第一反应是把缺失像素置为0。这个操作虽然可以让训练跑通但会引入一个隐患网络会学到“深度值为0的位置就是物体边界”而这个规律在真实场景中并不成立因为远距离区域的深度也可能被截断到0附近。我的处理方式是在输入层做两步。第一步缺失像素统一重设为一个“远离量程”的常量比如量程是8米缺失就设为8.5然后归一化时不会和其他正常深度混在一起第二步对缺失区域生成一个valid mask输入给Depth分支让网络明确知道哪些位置的深度不可信。实验下来这个组合比直接把缺失区域抠掉或者补成0都要稳定。训练时还可以做一种随机mask增强随机丢弃某个区域内的深度值强迫模型学习在部分深度缺失时仍然能保持分割质量。实测下来这种增强能显著提升模型在真实传感器上的鲁棒性因为开源数据集的深度质量普遍比自采数据好不做这个增强的话模型很容易出现过拟合。5.2 蒸馏信号不稳定怎么办特征蒸馏loss的数值曲线如果出现大幅震荡往往不是方法问题而是设备问题。最常见的原因是特征匹配层的数值量级不对齐一个分支的输出值域是[0,1]级别另一个是[-10,10]级别直接算MSE会让梯度被绝对值大的那个分支带偏。在投影层后加一个LayerNorm或BatchNorm强制把特征归一化到相同尺度震荡问题基本能解决大半。还有一个隐蔽问题就是两个分支的训练进度不一致。如果RGB分支已经收敛到不错的状态而Depth分支还在早期挣扎蒸馏loss会给Depth分支非常大的梯度甚至有把已经训练好的RGB特征也给带坏的情况。这就是为什么我前面建议在RGB主干相对稳定之前可以先冻结主干只训练融合模块和辅助分支等辅助分支有一定表达能力了再解冻主干联合训练。如果logits蒸馏的KL散度在训练中期开始变成一条水平线说明辅助分支已经学不到更多来自教师的信息了。这时候可以去检查温度系数是不是太小概率分布太尖锐梯度几乎只在真实类别那一维上有值其他维度全是0学不到软标签里的类别关系信息。5.3 工程效率与显存优化双backbone结构最直接的问题就是显存开销接近翻倍。我实际测试下来输入分辨率如果跑到512x512batch size设8单卡16G是基本门槛。优化手段有几种深度分支选用更轻量的结构不一定要和RGB共享同一个backboneDepth特征不需要那么高的语义抽象能力轻量网络足够在stage之间插入融合点而不是每个stage都融合融合点越多梯度路径越复杂显存占用也越多3个融合点是一个比较均衡的选择使用梯度累积等效扩大batch size能有效稳定BN层统计量多模态训练对batch size的敏感度比单模态更高batch太小BN跑偏的概率大幅增加。另一个工程细节是Depth和RGB的输入分辨率对齐问题。有些公开数据集的深度图和RGB图尺寸不一致DataLoader里必须做好resize和crop的统一。千万别用两套不同的resize参数否则后面所有工作都白做。写在最后我最早复现这套主模态引导方案时其实是持怀疑态度的觉得无非是给深度分支加了点约束能有多大区别。实际跑下来才知道多模态融合真正的问题往往不是出在融合模块本身而是两个分支的特征空间能不能顺畅对话。PrimKD不搞复杂的注意力花活用一个蒸馏约束把两个分支的表示对齐了后面所有融合模块都变得好训练了。那个“先对齐再放开”的训练策略是我在整个复现过程中收获最大的经验。反过来想这套“主模态蒸馏”的思路其实可以迁移到很多模态组合上RGB加LiDAR点云、RGB加事件相机、语言加视觉都存在类似的信息不对称和特征空间不匹配问题。只要你能找到那个信息密度最高的主模态再用蒸馏去引导辅助模态对齐大概率都能得到一组更好的融合效果。说到底多模态融合不应该是把所有输入一视同仁地塞进同一个网络而是搞清楚每个模态擅长什么再让它们各司其职。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
