【Uformer论文阅读｜CVPR 2022】：通用U型Transformer架构，重新定义图像修复任务

论文信息标题Uformer: A General U-Shaped Transformer for Image Restoration会议CVPR 2022单位中国科学技术大学、澳门大学、中国科学院大学代码https://github.com/ZhendongWang6/Uformer论文https://arxiv.org/pdf/2106.03106.pdf一、前言图像修复的痛点与Transformer的破局图像修复去噪、去模糊、去雨、去摩尔纹一直是低阶视觉的核心任务传统CNN擅长捕捉局部纹理但长距离依赖建模弱原生Transformer全局自注意力虽能抓长距离依赖但高分辨率图计算量爆炸且局部细节捕捉能力拉胯。Uformer直接给出两全方案U型架构LeWin Transformer块多尺度恢复调制器兼顾全局依赖与局部细节计算量还低在多个图像修复任务刷到SOTA。二、核心总览Uformer到底强在哪两大原创杀招LeWin Transformer块非重叠窗口自注意力降计算量深度卷积补局部细节解决Transformer“局部弱、计算贵”问题。多尺度恢复调制器轻量可学习空间偏置解码器多层校准特征细节恢复拉满几乎不增参量和计算量。三、网络整体架构U型Transformer的优雅设计3.1 整体流水线图1 Uformer整体架构输入退化图像→3×3卷积LeakyReLU提取低层特征。编码器4级层级结构每级若干LeWin块下采样4×4卷积步长2通道翻倍、尺寸减半。瓶颈层LeWin块低分辨率特征抓全局依赖。解码器4级结构每级上采样2×2转置卷积步长2跳跃连接LeWin块调制器。输出3×3卷积得残差图I′IRIIRI′IR得到修复图像。3.2 公式解析修复结果计算I′IRIIRI′IRIII输入退化图像尺寸3×H×W3×H×W3×H×WRRR网络输出残差图像尺寸3×H×W3×H×W3×H×WI′II′最终修复图像通俗解释网络不学直接输出清晰图只学退化与清晰图的差值训练更稳、细节更准。四、核心模块1LeWin Transformer块局部增强窗口Transformer4.1 设计初衷原生Transformer两大坑全局自注意力计算量O(H2W2C)O(H^2W^2C)O(H2W2C)高分辨率图跑不动。FFN纯全连接局部上下文捕捉弱图像修复需要邻域信息修复像素。LeWin块W-MSA窗口多头自注意力LeFF局部增强前馈网络残差层归一化标准配置。4.2 计算流程公式通俗解释Xl′W-MSA(LN(Xl−1))Xl−1X_l \text{W-MSA}(\text{LN}(X_{l-1})) X_{l-1}Xl′​W-MSA(LN(Xl−1​))Xl−1​XlLeFF(LN(Xl′))Xl′X_l \text{LeFF}(\text{LN}(X_l)) X_lXl​LeFF(LN(Xl′​))Xl′​Xl−1X_{l-1}Xl−1​第l−1l-1l−1块输出特征LN\text{LN}LN层归一化稳定训练W-MSA\text{W-MSA}W-MSA窗口自注意力抓长距离依赖LeFF\text{LeFF}LeFF局部增强前馈抓局部纹理Xl′、XlX_l、X_lXl′​、Xl​模块中间/最终输出通俗解释先做窗口注意力抓全局关系再做带卷积的前馈抓局部细节残差连接保证梯度不消失。4.3 W-MSA非重叠窗口自注意力图2 LeWin Transformer块结构特征图切分为M×MM×MM×M非重叠窗口窗口内做自注意力。计算量从O(H2W2C)O(H^2W^2C)O(H2W2C)降至O(M2HWC)O(M^2HWC)O(M2HWC)高分辨率图也能跑。加入相对位置编码BBB提升空间信息建模。4.4 注意力公式全解析Attention(Q,K,V)SoftMax(QKTdkB)V\text{Attention}(Q,K,V)\text{SoftMax}\left(\frac{QK^T}{\sqrt{d_k}}B\right)VAttention(Q,K,V)SoftMax(dk​​QKT​B)VQQQ查询矩阵表征当前特征KKK键矩阵表征所有特征VVV值矩阵特征真实信息dkd_kdk​单头注意力维度dkC/kd_kC/kdk​C/kkkk为头数BBB相对位置偏置可学习参数表通俗解释窗口内特征互相“注意力打分”加位置偏置保证空间关系不丢失计算快还准。4.5 LeFF局部增强前馈网络图3 LeFF局部增强前馈网络流程线性层升维→Reshape为2D特征图。3×3深度卷积抓局部邻域信息补Transformer局部弱短板。Reshape回1D→线性层降维→GELU激活。通俗解释给纯全连接的FFN加个小卷积让模型学会看像素邻居修复边缘、纹理更精准。五、核心模块2多尺度恢复调制器细节增强神器5.1 设计思路不同退化噪声、模糊、雨纹模式不同需要针对性特征校准。调制器可学习M×M×CM×M×CM×M×C张量解码器每个LeWin块的自注意力前作为共享偏置加入窗口特征。轻量仅加法操作参量/计算量几乎不增。通用适配去噪、去模糊、去雨等所有修复任务。5.2 效果可视化图4 调制器对去噪/去模糊的效果提升无调制器残留模糊/噪声细节模糊有调制器退化去除彻底文字、边缘清晰还原通俗解释给解码器每层加“细节滤镜”针对性修正特征修复效果肉眼可见提升。六、损失函数Charbonnier Lossℓ(I′,I^)∥I′−I^∥2ϵ2\ell(I,\hat{I})\sqrt{\|I-\hat{I}\|^2\epsilon^2}ℓ(I′,I^)∥I′−I^∥2ϵ2​I′II′模型输出修复图I^\hat{I}I^真实清晰图ϵ\epsilonϵ常数10−310^{-3}10−3防止根号内为0通俗解释比L2更鲁棒对异常值不敏感图像修复训练更稳定。七、核心代码实现PyTorch7.1 WindowAttention窗口自注意力importtorchimporttorch.nnasnnclassWindowAttention(nn.Module):def__init__(self,dim,window_size,num_heads):super().__init__()self.dimdim self.window_sizewindow_size self.num_headsnum_heads self.scale(dim//num_heads)**-0.5self.qkvnn.Linear(dim,dim*3)self.projnn.Linear(dim,dim)# 相对位置偏置表self.relative_position_bias_tablenn.Parameter(torch.zeros((2*window_size-1)**2,num_heads))defforward(self,x):B,N,Cx.shape qkvself.qkv(x).reshape(B,N,3,self.num_heads,C//self.num_heads)q,k,vqkv.unbind(2)qq*self.scale attn(q k.transpose(-2,-1))# 加相对位置偏置relative_position_biasself.relative_position_bias_table.view((self.window_size*self.window_size),(self.window_size*self.window_size),-1)attnattnrelative_position_bias.permute(2,0,1).contiguous()attnattn.softmax(dim-1)x(attn v).transpose(1,2).reshape(B,N,C)xself.proj(x)returnx7.2 LeFF局部增强前馈网络classLeFF(nn.Module):def__init__(self,dim,expand_ratio2):super().__init__()self.dimdim self.expand_dimdim*expand_ratio self.fc1nn.Linear(dim,self.expand_dim)self.dwconvnn.Conv2d(self.expand_dim,self.expand_dim,kernel_size3,padding1,groupsself.expand_dim)self.fc2nn.Linear(self.expand_dim,dim)self.actnn.GELU()defforward(self,x,H,W):B,N,Cx.shape xself.fc1(x)# Reshape为2D做深度卷积xx.transpose(1,2).view(B,self.expand_dim,H,W)xself.dwconv(x)xx.view(B,self.expand_dim,N).transpose(1,2)xself.act(x)xself.fc2(x)returnx7.3 LeWin Transformer块classLeWinBlock(nn.Module):def__init__(self,dim,window_size,num_heads):super().__init__()self.norm1nn.LayerNorm(dim)self.attnWindowAttention(dim,window_size,num_heads)self.norm2nn.LayerNorm(dim)self.leffLeFF(dim)defforward(self,x,H,W):# 窗口注意力残差xxself.attn(self.norm1(x))# LeFF残差xxself.leff(self.norm2(x),H,W)returnx八、实验结果全方位SOTA8.1 真实噪声去除表格1 出处原论文Table 1方法SIDD-PSNRSIDD-SSIMDND-PSNRDND-SSIMBM3D25.650.68534.510.851NBNet39.750.95939.890.955Uformer-B39.890.96039.980.955表格1 真实噪声去除结果Uformer-B在SIDD/DND均超NBNet参量更少、计算量更低。小模型Uformer-T性能超多数算法计算量最小兼顾速度与精度。8.2 运动去模糊表格2 出处原论文Table 2方法GoPro-PSNRRealBlur-R-PSNRRealBlur-J-PSNRMPRNet32.6635.9928.70Uformer-B32.9736.2229.06表格2 运动去模糊结果合成数据集GoPro、真实数据集RealBlur均超MPRNet泛化性极强。8.3 其他任务散焦去模糊DPD数据集超KPAC1.04dB。真实去雨SPAD数据集达47.84dB超前任3.74dB。去摩尔纹TIP18数据集超UNet2.79dB。九、消融实验验证每个模块的价值9.1 Transformer vs CNN表格3 出处原论文Table 5模型GMACs参数量PSNRUNet-B86.97G53.58M39.71Uformer-B89.46G50.88M39.89表格3 Transformer与CNN架构对比Uformer-B参量更少PSNR超UNet-B0.18dBLeWin块有效性拉满。9.2 局部增强位置表格4 出处原论文Table 6增强位置PSNR无增强39.74注意力中增强39.72FFN中增强39.77表格4 局部增强位置对比局部增强加在FFN效果最好加在注意力反而下降符合图像修复局部依赖逻辑。9.3 调制器效果表格5 出处原论文Table 7任务无调制器有调制器提升去模糊(GoPro)29.1129.570.46dB去噪(SIDD)39.8639.890.03dB去雨(SPAD)47.4347.840.41dB表格5 调制器消融结果调制器在所有任务均涨点计算/参量几乎无开销性价比拉满。十、总结Uformer是通用、高效、高性能的图像修复Transformer架构LeWin块窗口自注意力降计算深度卷积补局部解决Transformer两大痛点。多尺度调制器轻量校准特征细节拉满适配所有修复任务。U型架构层级特征跳跃连接多尺度信息融合拉满。去噪、去模糊、去雨、去摩尔纹全任务SOTA速度与精度双赢。