项目地址:
GitHub - swz30/Restormer: [CVPR 2022--Oral] Restormer: Efficient Transformer for High-Resolution Image Restoration. SOTA for motion deblurring, image deraining, denoising (Gaussian/real data), and defocus deblurring.
论文地址:
Restormer: Efficient Transformer for High-Resolution Image Restoration
应用:defocus deblurring(single-image and dual-pixel data), and image denoising(Gaussian grayscale/color denoising, and real image de-noising).
期刊: CVPR 2022
概述
和MIRnet一样,都来自于阿联酋的Inception Institute of Artificial Intelligence。这篇文章聚焦于解决Transformer应用在vision中的痛点。
Vision Transformers(ViT)把图像分成多个patch,学习他们之间的关系。所以相比于CNN,Vision Transformers的优点是更大的感受野,可以对图像内容有自适应的能力。
Thedistinguishing feature of these models is the strong ca-pability to learn long-range dependencies between imagepatch sequences and adaptability to given input content.
Vision Transformers的问题是SA模块的计算量与patch的分辨率是二次指数的关系,这就限制了Vision Transformers的实际应用。
一种做法是把self attention局限在局部区域,相当于加一个窗,Swin Transformer 就是这样的思路。但这样把信息集成局限在局部,削弱了Vision Transformers原有的优势,违背了初衷,不适用于图像复原任务。
所以这里选择了另外一条路:使用多头的SA+多尺度的结构。
结构
整体结构的特点:
1. 在卷积得到特征层之后,仍然是Unet的结构+skip connect,只不过每一层由Transformer Block构成。
2. u-net的deep feature再经过一个tansformer block进行refine,得到refined features。
3.refined features经过卷积层得到残差图像R,残差图R和原始图相加得到最终的复原图。
encoder-decoder
encoder-decoder共四层,每一层由tansformer block构成,由小到大表示四个尺度。在最后的level之后又经过transformer block,融合了encoder的low-level image features和decoder的 high-levelfeatures得到了deep feature。
encoder-decoder是对称的四层。它的输入是3x3卷积之后得到的low-level image features,使用表示,encoder-decoder的输出是deep features,使用 表示。
像u-net一样,encoder部分的四层的分辨率越来越低,但是通道层面越来越宽。decoder则逐步恢复分辨率。下采样和上采样分别使用pixel-unshuffle和pixel-shuffle实现。
在每一层中,为了辅助恢复的过程,encoder下采样之前的数据和decoder上采样后的数据会进行信息整合,使用的操作是通道层面的concatenation。为了避免通道数过多,会再使用1x1的卷积进行通道层面的压缩(压缩至一半)。
Transformer Block
MDTA(Multi-Dconv Head Transposed Attention)
MDTA作为transformer block的第一部分。
传统的SA在空域中计算点积,空域本身就是二维的,所以传统的SA的复杂度是。而在MDTA中是在通道之间计算点积,本质上是计算通道之间的协方差cross-covariance。
所以可以看到下图中的通道维度C被标红了:
除了通道方面的改进,可以看到Q、K、V都经过了两层卷积,分别是1x1和3x3的深度可分离卷积,并且是bias-free的。MDTA在计算协方差之前,使用深度可分离卷积depth-wise convolutions,可以强调局部语义信息。
得到Q、K、V之后,还会进行reshape操作,把H和W的乘积看作一个维度,通道看作另外一维。因为Q和K在softmax之前会进行内积的计算,reshape之后再内积的结果是CxC的,这就达到了通道之间协方差的表示,称之为transposed-attention map A。
transposed-attention map A和V再做内积,达到注意力机制的目的。此时得到的仍然是HWxC大小,再进行reshape恢复成三维度的数据。
再经过1x1卷积得到残差,和输入求和,得到最终的输出。
GDFN(Gated-Dconv Feed-Forward Network)
GDFN作为transformer block的第二部分。MDTA的作用是特征提取,GDFN的作用是特征转换。
特征转换通常使用regular feed-forward network(FN),使用两次1x1卷积,先扩展通道,再压缩通道。两次1x1卷积还能引入非线性。
GDFN仍然使用了1x1卷积和3x3深度可分离卷积结合的方式,并且使用了门逻辑来操作对应像素。具体做法是分成两路:一路经过GELU non-linearity激活器,激活的结果作为权重,施加到另外一路上去。对应像素按权重求和,得到的结果
GDFN的本质是控制信息流。因为每一层的transformer都有GDFN,所以不同层可以专注于提取不同的互补的特征。
总结
先看看疗效,Restormer 在效果和速度上都能找到优势:
0.3 dB and 0.25 dB over the previous best CNN method MIRNet and Transformer model Uformer;
compared to SwinIR , Restormer has 3.14× fewer FLOPs and runs 13× faster
restomer中的两个重要改进其实都是多头基础上进行的:不同头直接经过门响应或者内积得到权重,对另外的分支进行加权。
在后面的文章中,可以看到NAFnet对此进行了进一步的简化:把多头使用多个通道代替,并且因为对应乘积本身有非线性,激活层甚至都可以去除掉。
reference:
论文及代码详解——Restormer-CSDN博客