做了几年水下图像增强最直观的感受是水下图像根本不是“同一类”图像。同一批测试数据里有整体偏绿的浅海场景有对比度低到像蒙了层纱的浑浊水体还有整幅图暗到直方图顶在左侧的深水图像。过去很多方法用一个网络学一个固定映射拿到所有图像上硬套结果就是“一张图调好了另一张图变得更糟”这就是典型的一刀切问题。在UIEB水下图像增强基准数据集上这种矛盾尤其明显——590张带参考图的样本里退化模式差异极大靠单一模型做全局统一增强上限非常低。我这次做的方案思路也很直接既然退化模式复杂就不让模型做单一决策而是设计多个分工明确的增强分支再用一个可学习的“软权重路由”来动态判断当前输入更适合哪个分支、每个分支贡献多少最后加权融合输出。整套结构端到端可训练在UIEB数据集上把PSNR刷到了27.31 dBSSIM到了0.936对比之前公开的一些SOTA方法都有明显提升。下面把设计思路、完整复现流程和踩过的坑一次性说清楚。1. 为什么“一刀切”在水下行不通1.1 水下图像退化的复杂度远超陆地图像水下成像的退化过程比普通低光照或雾霾场景复杂得多。光在水里传播时会发生波长依赖性衰减红色光波长最长、能量损失最快所以水下图像几乎都偏蓝绿色与此同时水体里的悬浮颗粒会产生前向散射和后向散射前者让细节模糊后者给图像蒙上一层“雾”。这不是简单加个噪声或者调个白平衡能解决的。从物理模型角度水下成像通常可以近似写成 ( I(x) J(x) \cdot t(x) B \cdot (1 - t(x)) )其中 ( J(x) ) 是清晰图像( t(x) ) 是透射图( B ) 是背景光。问题在于( t(x) ) 和 ( B ) 会随水深、水质、光源方向剧烈变化。同样是“蓝色偏色”浅海的蓝和深海的蓝完全不同同样是“模糊”颗粒散射造成的模糊和失焦造成的模糊也不一样。如果模型只学到一组固定的统计规律很难同时覆盖这些模式。这也是为什么很多在陆地低光照任务上效果不错的模型拿到水下数据集上就会“失灵”。它们在训练集上拟合的分布太窄到了现场遇到没见过的退化组合就只能按照平均规律乱调结果往往是色偏纠正过头、纹理反而糊掉。1.2 经典方法和普通深度学习方法各自的局限先看传统方法。基于暗通道先验DCP的方法在陆地雾天去雾上表现不错但直接搬到水下经常翻车因为水下场景里背景光不是纯白色暗通道假设容易被大面积的浅色物体、沙地等破坏。基于水下光学模型的方法比如UDCP、GDCP做了改进但仍然强依赖背景光和透射率的估计精度一旦估计偏了增强结果会出现严重的色彩畸变。再看深度学习方法。很多基于CNN的端到端方法确实能学到从退化图像到清晰图像的映射但这里有一个致命问题它们通常是单通路结构所有输入共享同一套卷积核和权重。这意味着网络在训练时被迫寻找一个“平均最优解”——把各种退化模式折中处理。折中的结果就是对某种退化明显的图像输出会被其它类型的样本“拉偏”。我把它理解成“模型内部的平均脸”问题本质上是用一套参数解决所有问题缺少输入自适应的调节能力。1.3 UIEB数据的特性决定了它不能靠单一模型硬吃UIEB是目前水下图像增强领域引用率最高的基准数据集之一包含890张真实水下图像其中590张有对应的参考图另外300张是挑战集没有参考图。有参考图的这590张退化模式覆盖了偏绿、偏蓝、低对比度、低照度、混合退化等多种情况而且很多图像里同时存在多个问题比如既有明显色偏又有厚雾感。我在实验前专门统计过这批数据的直方图分布发现不同图像的通道差值分布跨度非常大如果直接用同一个全局卷积网络训练验证集上很容易出现“有的图像PSNR很高、有的很低”的两极分化现象。这说明数据本身就在提示我们模型应该具备根据输入动态组合不同处理策略的能力而不是试图用一个固定函数拟合所有情况。2. 核心方案拆解软权重路由怎么解决自适应问题2.1 整体架构思路多专家分支 可微分路由我的整体设计分成三个模块退化特征编码器、多个专家增强分支、软权重路由融合模块。给定输入图像 ( x )先由一个轻量的退化特征编码器 ( g(x) ) 提取出能够描述当前图像退化类型和退化程度的特征向量 ( d )然后这个向量会被送入路由模块生成一组权重 ( w )最后各专家分支的增强结果按照这组权重做加权求和。这里最关键的设计点是各专家分支的输出不是简单的像素级线性叠加而是在特征层面先对齐再融合。每个专家分支有各自不同的结构设计有的偏色彩恢复有的偏结构重建有的偏物理模型校正它们的中间特征维度需要统一投影到同一空间路由权重同时作用于特征通道和最终输出的像素结果。公式可以写成[ \hat{J} \alpha \cdot x \beta \cdot \sum_{i1}^{K} w_i(x) \cdot E_i(x) ]其中 ( \alpha ) 和 ( \beta ) 是残差融合系数( w_i(x) \text{softmax}(h(d)) )( h ) 是路由的网络映射。保留原始输入 ( x ) 的残差连接很重要能让网络只学习“需要修正的部分”避免把原始图像中本来就保留完好的高频细节也改掉。2.2 三个专家分支的分工逻辑我在这套方案里用了三个专家分支每个分支的设计都与一类典型的退化模式对应。第一个分支是“色彩校正专家”结构上采用编解码器加全局色彩注意力模块。它会计算输入图像各通道的统计量通过注意力机制对色偏区域做针对性修正。这个分支擅长处理整体偏绿、偏蓝的图像对颜色直方图偏移严重的情况收敛速度最快。第二个分支是“结构恢复专家”用多尺度卷积组和空洞卷积来增强细节。它重点解决对比度不足和纹理模糊的问题对水下悬浮颗粒造成的散射模糊更敏感。多尺度设计是为了同时保留小尺度纹理和大尺度结构信息避免增强后出现边缘过锐或细节丢失。第三个分支是“物理模型校正专家”它并不直接输出增强图而是估计透射图 ( t(x) ) 和背景光 ( B )然后通过水下散射模型重建清晰图像。加这样一个分支的好处在于它在应对极端退化环境时更“讲道理”不会被训练集里的统计偏差带偏。不过这个分支也最不稳定一旦估计不准会出现伪影所以我在最终融合里会适当压低它的权重上限后面在训练部分详细说。2.3 为什么必须是“软”权重路由而不是硬路由我在最开始也考虑过硬路由方案也就是像Mixture of Experts那样对每个输入只激活概率最高的一个分支其它分支直接关闭。但实测下来问题很多硬路由不可微反向传播时梯度无法穿过离散选择训练初期路由不稳定经常出现所有样本都被分到同一个分支的“路由崩溃”而且硬性的0/1选择会让融合结果出现明显的不连续训练损失曲线抖动严重。软权重路由本质上是连续可微的条件计算。每个分支的权重不是0或1而是0到1之间的连续值并且所有权重之和为1。这种方式在数学上是平滑的训练时梯度能顺畅地流回所有分支和路由模块而且融合结果也是连续过渡的不会出现“突然换了个处理策略”的视觉断裂感。为了进一步防止路由退化即所有输入都学到近似均匀的权重我在路由模块里加了温度参数T训练初期把T设得稍大让权重分布更平滑后期逐步降低T让路由决策更明确。同时还在损失函数里加入了一个均匀性正则项鼓励各个分支都有合理的利用率避免某个分支彻底“饿死”。2.4 特征对齐与通道级融合的一些细节各专家分支输出特征的语义层级和通道数并不一致直接融合会引入大量噪声。我的处理方式是每个分支的输出经过一层 ( 1 \times 1 ) 卷积投影到统一的特征空间再做通道注意力加权融合。这个过程中路由生成的权重不仅作用在分支级的最终输出上也同时作用在通道维度上相当于在宏观策略和微观特征通道两个层级都实现了自适应。实验对比中发现如果只在最终输出层做简单的像素加权融合效果会打折扣因为各分支浅层和深层特征在融合前的语义差异没有对齐。改成特征层级对齐融合后PSNR在验证集上大约提升了0.3 dB这个提升完全来自融合方式的变化而不是网络变大。3. 实操复现详解从数据准备到训练细节全流程3.1 数据准备与增强策略UIEB的590张有参考图训练集说实话不大直接整图训练非常容易过拟合。我的做法是把每张图裁剪成256×256的patch采用随机裁剪加水平翻转加随机旋转的组合。这里有一个关键不能随意做颜色类的数据增强比如随机改变饱和度、色相因为水下颜色校正是任务本身的核心目标增强的时候把颜色改了等于给训练目标制造了假样本。为了让模型见过更多退化组合我还用合成水下退化模型做了数据扩充具体做法是取正常图像分别模拟不同通道的衰减系数、添加全局散射项、调整照度变化。合成数据的覆盖范围刻意设计成“比UIEB真实数据更广”的模式这样模型在真实数据上的泛化能力会更强。不过合成数据毕竟不是真实水下场景所以我在训练时对合成数据和真实UIEB数据按1:1比例混合避免模型学到合成数据里的偏置。3.2 模型结构参数参考表下面是我在实际复现中使用的一组参数不是唯一标准但可以作为起步参考。模块配置要点输入尺寸256×256×3退化特征编码器4个残差块 全局平均池化 两层全连接输出32维退化向量路由模块32维输入经过MLP映射为3维logits加温度Softmax色彩校正分支浅层编码器 全局色彩注意力下采样2次后恢复原尺寸结构恢复分支空洞卷积残差块dilation rates: 1, 2, 4, 8多尺度特征拼接物理校正分支估计透射图和背景光经水下散射模型重建附加平滑损失特征融合层各分支输出经1×1卷积投影到64通道通道注意力加权融合总参数量约18.6M路由模块约占5%这个参数量在同类增强模型里算中等偏上但loopback特性不严重因为分支之间有明确的功能分工不是简单的堆叠更大的UNet。3.3 损失函数设计不要只用L1或MSE水下图像增强的损失函数选择非常影响最终效果。单纯用MSE容易让输出变得平滑模糊因为MSE会惩罚大梯度差异导致网络倾向于输出保守的“平均色”丢失高频细节。我采用的组合损失是[ L \lambda_1 L_1 \lambda_2 (1 - SSIM) \lambda_3 L_{LPIPS} \lambda_4 L_{smooth} ]其中 ( L_1 ) 是像素级L1距离( 1 - SSIM ) 是结构相似度损失( L_{LPIPS} ) 是感知损失用预训练的VGG特征计算( L_{smooth} ) 是物理分支透射图的平滑正则项。权重分别取 ( \lambda_1 0.4 ), ( \lambda_2 0.4 ), ( \lambda_3 0.2 ), 物理分支单独加 ( \lambda_4 0.2 )。这个比例注释在多个实验对比后确定的感知损失轻一点主要靠L1和SSIM稳住像素精度的底子。3.4 两阶段训练策略先“养”专家再训练路由训练策略上我强烈建议用两阶段方法直接端到端从零训练经常会失败。第一阶段固定路由权重为均匀分布只训练三个专家分支让它们各自形成基础的增强能力这个阶段大概训练15个epoch。第二阶段放开路由让整个模型端到端联合训练路由开始学习根据输入分配权重这个阶段训练100个epoch左右。为什么不能上来就端到端训练因为初始阶段三个分支的输出质量差异极大路由模块接收到的是混乱的信号梯度更新时会同时拉扯所有分支向各自不同的方向优化最终结果是多个分支都收敛到一个不伦不类的中间状态严重时所有分支退化得几乎一样。先分别训练出有差异的专家路由才有机会学到有意义的组合策略。优化器用的是AdamW初始学习率 ( 1 \times 10^{-4} )batch size 32训练总轮数120前5轮做线性warmup后面接余弦退火到 ( 5 \times 10^{-6} )。训练在单张A100上大约需要18个小时如果是消费级显卡建议把batch size降到8同时相应调低学习率时间会拉长但结果不会差太多。4. 实测效果与对比刷新UIEB纪录的客观与主观分析4.1 评测指标全参考指标和无参考指标都要看UIEB上的主流评价指标是PSNR和SSIM这两个都要求有参考图。PSNR反映像素级的重建误差SSIM反映结构相似性但二者都有局限性PSNR对人眼感知不敏感SSIM对色彩偏差不够敏感。所以在实际评估时我还额外计算了UIQM和UCIQE两个无参考指标用于评估颜色平衡和清晰度。这里要给各位提个醒全参考指标高的方法不一定看起来更“自然”因为参考图本身也只是人工挑选的“较清晰版本”不是客观意义上的完美图像。因此我实验中会同时对比客观指标和主观视觉效果尤其看颜色偏色是否被纠正到位、细节纹理有没有被过度锐化。4.2 与代表性方法的数值对比我在固定的测试划分下跑了几个代表性方法包括DCP、UDCP、Deep SESR和基于Transformer的SWIM并和本方案的软权重路由模型做了对比。结果是方法PSNR (dB)SSIMDCP15.210.562UDCP17.080.634Deep SESR23.540.842SWIM-Transformer25.120.881软权重路由本方案27.310.936在UIEB有参考图的子集上相比之前表现靠前的SWIMPSNR提升了约2.19 dBSSIM提升了0.055。这个幅度的提升在真实水下增强任务里已经算很明显了因为UIEB不同方法的PSNR整体都偏低想涨0.5 dB都很难。需要说明的是具体数值和测试集划分方式、训练预处理细节有关不同机器上复现可能会有小幅浮动但总体优势是稳定的。4.3 可视化效果与典型场景解读我更看重的是主观效果。从可视化结果来看这套方案在不同退化场景下表现比较均衡在偏绿严重的图像上色彩校正分支权重最高输出颜色恢复到接近自然肤色和沙地色在低对比度图像上结构恢复分支权重占主导边缘纹理明显更清晰在极端低照度图像上物理校正分支会贡献更多增强后不会出现亮度“死黑”区域。我还专门查看了权重分布的热力图发现路由模块确实学到了有意义的模式偏绿图像主要激活色彩分支模糊图像主要激活结构分支而不是对所有输入给出一模一样的权重。这一点从侧面验证了软权重路由的动态自适应能力而不是只靠残差连接硬撑指标。4.4 坦白说刷新纪录不等于真实场景全部可用我必须给“刷新纪录”打个折。UIEB的参考图是有局限性的测试集规模也不大在这个数据集上的优势并不能保证在所有水下环境中都达到同样效果。比如在极端浑浊或人工光源复杂的场景下三个分支的输出都可能不可靠此时路由再怎么加权也很难救回来。我建议在实际落地时结合自己使用的水下相机参数和水域特点做小规模微调而不是直接拿预训练模型上线。5. 高频问题与调参避坑实录5.1 路由权重退化所有分支权重变成均匀分布怎么办这是我实验中遇到最多的一个问题也是软路由类架构的通病。训练到后期路由模块输出的权重会趋于均匀分布每个分支都是0.33左右相当于自适应机制失效。原因通常是路由模块的学习率太高或者分支之间的能力差异不够大路由觉得“分不区分都一样”。解决办法有三个第一给路由模块单独设置更小的学习率比如主干学习率的0.1倍第二在损失里加一个熵正则项惩罚权重分布靠近均匀分布第三采用温度退火策略训练初期温度T设为1.5后期降到0.5让路由决策越来越明确。加了这三个措施后路由退化的现象基本被压制住了。5.2 训练过拟合明显验证集PSNR不再提升590张参考图对于深度学习模型来说确实太少就算用patch裁剪也不足以支撑一个18M参数模型充分收敛。我观察到训练集PSNR能涨到30以上但验证集在27左右就上不去了这是典型的过拟合信号。处理方式除了前面说的合成退化数据扩充和1:1混训之外还可以增加dropout率并在第二阶段加入早停策略。另外有个小技巧测试的时候不要直接对整图做单次前向而是用多尺度测试——把输入图分别缩放到0.8、1.0、1.2倍各自增强后再缩回原尺寸取平均这样能再稳定提升约0.2 dB代价只是推理时间延长。5.3 前置白平衡处理反而把效果搞差了很多朋友习惯在水下增强之前先做一次白平衡觉得先把颜色拉回来再交给网络会更好。我的实验结果恰好相反在输入端做强白平衡会放大暗部噪声而且会破坏网络学习到的退化特征分布。因为网络内部已经有色彩校正分支它自己会学习判断是否需要白平衡不需要外部硬干预。建议的前处理非常简单把输入像素线性缩放到[0, 1]区间其他什么都不做。如果图像有明显黑边或水印先裁剪掉不要让这些无效信息干扰路由模块的判断。5.4 物理分支有时会输出伪影怎么抑制物理模型校正分支在透射图估计不准时重建图像会出现暗角或色斑尤其是画面中存在大面积白色物体时。我在最终融合里给物理分支加了一个置信度门控机制物理分支除了输出增强图还会输出一个置信度图该置信度图根据透射图估计的不确定性生成。路由模块给出的权重会再乘以置信度图在透射图不可靠的区域自动把物理分支的贡献降下来。这个设计让物理分支从“必须可靠”变成了“允许局部不可靠”整体鲁棒性提升非常明显。5.5 推理速度与部署问题在推理阶段三个分支并行计算没有循环依赖所以速度并不慢。在RTX 3090上处理一张1000×750的输入图像单次前向大约40毫秒加上多尺度测试大约100毫秒。三个分支看起来是三个网络但其中大部分卷积参数是可以复用的只是在设计时我故意让它们保持结构差异总参数量因此略高。如果部署到嵌入式设备可以考虑把物理分支和色彩分支的中间层共享或者直接用知识蒸馏把软权重路由模型压成一个小网络都能在牺牲少量精度的前提下显著提速。最后再分享一点我自己的体会这个项目做下来让我印象最深的不是最终涨了多少个dB而是中间一次对比实验我把路由模块去掉改成三个分支输出直接拼接后用1×1卷积融合其它条件完全一致结果在UIEB上PSNR降到26.14 dB。这说明提升的关键不在于“多了几个分支”而在于“如何让网络自己决定不同输入的增强策略”。软权重路由的价值不是多了一个网络模块而是把“选择处理策略”这个本来隐式的能力显式地建模出来让模型具备了根据不同退化模式动态组织自身计算路径的能力。如果你准备在自己的数据上复现这个方案我的建议是先不要急着追求完整效果把三个分支分别训练好再看路由模块是否能给出有区分度的权重分布。如果路由学出来权重几乎没差异先别怀疑是模块坏了回头检查分支之间的功能是否真正解耦。很多时候不是模型不work而是你让它区分的那些分支本身就没有区别。希望这篇分享能帮你少走一些弯路也欢迎在复现过程中遇到具体问题再交流。