049、多任务联合学习:超分与去噪、去模糊的协同优化框架
049、多任务联合学习超分与去噪、去模糊的协同优化框架上周调一个真实场景的超分模型输入是监控摄像头拍的夜间画面结果模型把噪点当纹理给放大了输出图像比低分辨率原图还难看。更离谱的是同一套模型在公开测试集上PSNR涨了0.3dB一到真实场景就崩。后来发现问题出在训练数据太“干净”——全是理想条件下的清晰图像而真实场景里噪声、模糊、下采样是同时存在的。这就是典型的“单任务超分”的坑你只教模型怎么放大没教它怎么处理退化。退化模型不是你想的那样很多论文里退化模型写成y (x * k) ↓s n其中k是模糊核↓s是下采样n是噪声。但实际场景中模糊和噪声是耦合的——模糊会改变噪声的分布噪声会掩盖模糊的细节。更麻烦的是不同退化参数之间还有交互效应比如高斯噪声和运动模糊同时存在时模型往往会把模糊边缘误判为噪声导致去噪过度而丢失结构。我踩过最深的坑是单独训练一个去噪网络和一个超分网络然后串起来用。结果去噪网络把高频细节抹掉了超分网络只能“脑补”出错误的纹理。后来改成联合训练损失函数里同时包含去噪损失和超分损失但梯度冲突严重——去噪任务希望平滑超分任务希望锐化两个梯度方向几乎相反。协同优化的核心共享编码器任务特定解码器我们团队最终采用的方案是一个共享的编码器提取多尺度特征然后分叉出三个解码器——去噪分支、去模糊分支、超分分支。关键点在于共享编码器要学到“退化无关”的表示也就是不管输入是带噪声还是带模糊编码器输出的特征应该聚焦在图像内容本身。这里有个trick在编码器中间层加入退化感知模块。具体做法是在特征图上做通道注意力但注意力权重由退化类型动态生成。比如输入是噪声主导注意力就放大去噪分支需要的特征通道输入是模糊主导就放大去模糊分支需要的边缘特征。这个模块用一个小型分类网络来预测退化类型然后生成注意力掩码。代码实现时注意退化感知模块的输入不能是原始图像而应该是编码器中间层的特征图。否则分类网络会学到图像内容信息导致注意力被内容主导而不是退化类型主导。别这样写# 错误示范用原始图像预测退化类型deg_typedeg_classifier(lr_image)# 会学到内容信息attentiongenerate_attention(deg_type)正确做法是# 正确用编码器特征预测退化类型shared_featencoder(lr_image)deg_typedeg_classifier(shared_feat.detach())# 这里detach防止梯度反传到编码器attentiongenerate_attention(deg_type)shared_featshared_feat*attention# 这里踩过坑attention要广播到所有通道损失函数的博弈梯度调和策略多任务学习的核心难题是梯度冲突。我们试过直接加权求和L_total λ1 * L_sr λ2 * L_denoise λ3 * L_deblur结果λ怎么调都不收敛。后来改用梯度调和Gradient Harmonization思路是计算每个任务的梯度向量如果两个梯度夹角大于90度即冲突就把其中一个梯度投影到另一个的法平面上。实现起来不复杂但要注意数值稳定性。这里给个简化版代码defharmonize_grads(grads_dict):# grads_dict: {sr: grad_sr, denoise: grad_denoise, deblur: grad_deblur}taskslist(grads_dict.keys())# 以超分梯度为基准base_gradgrads_dict[sr]fortaskintasks[1:]:ggrads_dict[task]dot(base_grad*g).sum()ifdot0:# 冲突# 将g投影到base_grad的法平面proj(dot/(base_grad.norm()**21e-8))*base_grad grads_dict[task]g-projreturngrads_dict注意这里1e-8是为了防止除零但实际训练中如果base_grad的范数太小投影会不稳定。我后来改成用余弦相似度做阈值只有当cos -0.3时才做投影避免过度修正。训练策略课程学习动态权重直接端到端训练容易陷入局部最优。我们的做法是分阶段第一阶段冻结共享编码器单独训练三个解码器。每个解码器用对应的退化数据训练比如去噪分支只用带噪声的图像超分分支只用清晰的下采样图像。这个阶段让每个分支学会自己的“专业能力”。第二阶段解冻编码器用联合损失微调。但此时三个任务的损失权重不是固定的而是根据验证集上的性能动态调整。具体做法是每个epoch结束后计算三个任务在验证集上的PSNR如果某个任务PSNR下降就增大它的损失权重。这里有个经验去噪任务和去模糊任务的权重不要同时调整否则模型会震荡。我通常先固定去噪权重调去模糊和超分的比例等稳定后再微调去噪权重。数据增强的隐藏技巧多任务联合训练对数据多样性要求极高。我们用的数据增强策略是对同一张高清图像随机生成不同的退化组合。比如模糊核高斯模糊σ∈[0.5, 3.0]或运动模糊长度∈[5, 20]像素噪声高斯噪声σ∈[0, 25]或泊松噪声λ∈[10, 50]下采样双三次插值尺度因子∈[2, 4]但要注意退化参数不能随机组合否则会出现“强模糊强噪声”这种极端情况模型学不到有用信息。我们的做法是模糊程度和噪声强度负相关——模糊越强噪声越弱反之亦然。这符合真实场景运动模糊通常发生在光线充足时噪声小而低光照下噪声大但模糊小。推理时的退化自适应训练好的模型在推理时输入图像可能包含未知的退化组合。我们加了一个退化估计模块先用一个小网络预测输入图像的模糊核大小和噪声标准差然后动态调整三个分支的融合权重。比如预测到噪声很大就放大去噪分支的输出权重预测到模糊严重就放大去模糊分支的贡献。这个模块在训练时是端到端学习的但推理时要注意退化估计网络的输入应该是原始低分辨率图像而不是经过任何预处理的版本。否则估计结果会偏差。个人经验总结别迷信公开数据集DIV2K、Flickr2K这些数据集太干净了训练出来的模型在真实场景大概率翻车。建议自己收集一些监控、手机拍摄的真实低质量图像用人工标注的方式做测试集。梯度调和比权重搜索更有效我花了两周调λ不如花两天实现梯度调和。而且调和后的训练更稳定不需要频繁调整超参数。共享编码器的深度要适中太浅5层学不到高层次特征太深20层会导致梯度消失尤其是去噪分支的梯度很难传到浅层。我常用的配置是12层残差块每层64通道。去噪和去模糊不要同时训练这两个任务在低层特征上有冲突。建议先训练去模糊分支等它收敛后再加入去噪分支。或者用交替训练一个epoch训练去模糊下一个epoch训练去噪。验证集要包含退化组合的极端情况比如只有噪声没有模糊、只有模糊没有噪声、两者都有。否则模型会在某个退化类型上过拟合。最后说一句多任务联合学习不是银弹。如果你的应用场景退化类型单一比如只有高斯噪声那单任务模型加上好的数据增强效果可能更好。只有在退化类型复杂且未知时联合学习才值得投入。