[论文解读] Uformer: A General U-Shaped Transformer for Image Restoration

Zhendong Wang, Xiaodong Cun+4 more|arXiv (Cornell University)|Jun 6, 2021

Advanced Image Processing Techniques参考文献 76被引用 109

一句话总结

Uformer 引入了一个带有局部增强窗口 LeWin 模块的U形 Transformer，以及一个轻量级多尺度恢复调制器，在降噪、去模糊、去散焦去模糊和除雾方面实现了最先进的结果，同时保持高效计算。

ABSTRACT

In this paper, we present Uformer, an effective and efficient Transformer-based architecture for image restoration, in which we build a hierarchical encoder-decoder network using the Transformer block. In Uformer, there are two core designs. First, we introduce a novel locally-enhanced window (LeWin) Transformer block, which performs nonoverlapping window-based self-attention instead of global self-attention. It significantly reduces the computational complexity on high resolution feature map while capturing local context. Second, we propose a learnable multi-scale restoration modulator in the form of a multi-scale spatial bias to adjust features in multiple layers of the Uformer decoder. Our modulator demonstrates superior capability for restoring details for various image restoration tasks while introducing marginal extra parameters and computational cost. Powered by these two designs, Uformer enjoys a high capability for capturing both local and global dependencies for image restoration. To evaluate our approach, extensive experiments are conducted on several image restoration tasks, including image denoising, motion deblurring, defocus deblurring and deraining. Without bells and whistles, our Uformer achieves superior or comparable performance compared with the state-of-the-art algorithms. The code and models are available at https://github.com/ZhendongWang6/Uformer.

研究动机与目标

• 动员在图像恢复中超越传统卷积网络，进行有效长程依赖建模的需求。
• 提出一种适用于多尺度图像恢复任务的通用U形 Transformer 架构。
• 开发高效的 LeWin Transformer 块，以在局部细节和全局上下文之间取得平衡。
• 引入一个轻量级的多尺度恢复调制器，以在各尺度上增强细节恢复。
• 在降噪、去模糊、散焦模糊和除雨数据集上展示最先进或具有竞争力的性能。

提出的方法

• 提出一种分层的UNet样编码-解码器，带跳跃连接，其中卷积被 LeWin Transformer 块所替代。
• 引入局部增强窗口（LeWin）Transformer 块，结合非重叠窗口自注意力（W-MSA）和带深度卷积的局部增强前馈网络（LeFF）。
• 对自注意力使用非重叠的 MxM 窗口，将复杂度从 O(H^2W^2C) 降低到 O(M^2HW C)。
• 将多尺度恢复调制器作为可学习的基于窗口的偏置加入解码器特征，以在各尺度上调整表示以适应恢复。

实验结果

研究问题

• RQ1 基于 Transformer 的U形架构，具备局部窗口自注意力和局部上下文 FFN，是否能有效同时捕捉局部细节与长程依赖用于图像恢复？
• RQ2 轻量级的多尺度恢复调制器是否在不显著增加计算开销的情况下，提升对多种降解类型的恢复质量？
• RQ3 在降噪、去模糊和除雨任务中，LeWin 块相对于传统CNN或全局注意力 Transformer 的性能与效率权衡是什么？

主要发现

• Uformer-B 在 SIDD 上达到 39.89 dB PSNR，在 DND 上达到 39.98 dB PSNR，超越了这些真实噪声数据集的前沿方法。
• 在运动去模糊上，Uformer 在 GoPro、RealBlur-R/J 和 HIDE 数据集上实现了最先进或具有竞争力的结果。
• 在散焦模糊上，Uformer 在 DPD 上的 PSNR 提升高达 1.87 dB，SSIM 也更高。
• 在真实降雨去除（SPAD）上，Uformer-B 达到 47.84 dB PSNR 和 0.9925 SSIM，比之前最好结果在 PSNR 上提高 3.74 dB。
• 消融实验显示 LeWin 块优于原生 UNet 变体，局部增强 FFN 有助于性能提升，调制器在 SPAD 尤其带来额外增益。
• 所提出的调制器在去模糊方面有显著提升（0.46 dB），在降噪和除雨任务中也有增益。