（2-论文精读、代码分析）Restormer: Efficient Transformer for High-Resolution Image Restoration CVPR2022
目录
1、Motivation
2、Contribution
3、Overall Structure
4、Method
4.1 MDTA
4.2 GDFN
4.3 TransformerBlock
4.4 Progressive Learning
4.5 测试实例
5、Experiments
5.1 Image deraining Results
5.2 Single-image Motion Deblurring Results
5.3 Defocus Deblurring Results
5.4 Image Denoising Results
5.5 Ablation Studies
6、Summary
[（1-复现使用手册）Restormer: Efficient Transformer for High-Resolution Image Restoration CVPR2022
( "（1-复现使用手册）Restormer: Efficient Transformer for High-Resolution Image Restoration CVPR2022")
这个论文的主要想法是将 Transformer 模型应用到图像修复中，不过和一般的VIT模型不同的是，没有使用 patch 级的特征进行 token 建模，可以理解为每个像素就是一个 token。
VIT 将输入图像划分为多个小 patch,每个 patch 被视为一个"token"，这种 patch-based token 建模方式,可以让 Transformer 学习到 patch 之间的相互关系,从而捕获图像中的全局信息。
Restormer 将每个像素视为一个 token，直接对图像中的每个像素进行建模和处理,而不需要先将图像划分为 patch，可以更好地捕获像素级的细节信息,而不会丢失局部信息。
Transformer block 主要包括两部分：一部分是self-attention 的计算，另一部分是 FFN 部分。作者也很自然的想到了改进这两个部分，提出了两个模块，分别是：
Multi-Dconv Head Transposed Attention，主要改进 self-attention 部分
Gated-Dconv Feed-Forward Network，主要改进FFN部分
1、Motivation
卷积神经网络(CNN)在图像修复任务上表现良好,主要原因是其能从大规模数据中学习到有效的图像先验。但CNN存在两个主要问题: 1) 卷积操作的感受野有限,无法建模长距离像素依赖关系; 2) 卷积滤波器在推理时权重是固定的,无法根据输入自适应调整。
自注意力机制和Transformer模型可以有效建模长距离依赖关系，在自然语言和高层视觉任务上取得了优异表现。但Transformer的计算复杂度随空间分辨率呈二次增长，无法应用于高分辨率图像修复任务。
在这项工作中，我们提出了一个有效的Transformer模型，通过在构建模块（多头注意力和前馈网络）中进行几个关键设计，使其可以捕获长距离像素交互，同时仍然适用于大图像。
我们的模型名为Restoration Transformer（Restormer），在多个图像恢复任务上实现了最先进的结果，包括图像去噪，单图像运动去模糊，散焦去模糊（单图像和双像素数据）和图像去噪（高斯灰度/彩色去噪和真实的图像去噪）。
2、Contribution
论文的主要创新点是将经典 Transformer 中的 MSA 和 FFN 都进行改进，并且采用 Encoder-Decoder 架构，涉及的上采样操作使用 nn.PixelShuffle() 来实现，涉及的下采样操作使用 nn.PixelUnshuffle() 来实现，整体文章的脉络非常清晰。
（1）提出了Restormer ，一个编码器-解码器Transformer，用于在高分辨率图像上进行多尺度局部-全局表示学习，而不将其分解为局部窗口，从而利用远程图像上下文。
（2）提出了一个多Dconv头转置注意（MDTA）模块 ，能够聚合本地和非本地像素的相互作用，是足够有效的处理高分辨率图像。
（3）一种新的门控Dconv前馈网络（GDFN） ，其执行受控特征变换，即抑制信息量较少的特征，并且仅允许有用的信息进一步通过网络分级结构。
主要目标是开发一个高效的Transformer模型，可以处理高分辨率的图像恢复任务。为了缓解计算瓶颈，将关键设计引入多头SA层和多尺度分层模块，其计算要求低于单尺度网络。
3、Overall Structure
论文的总体框架如下图所示，包括四层对称的 encoder-decoer 结构，还有一个用于 refinement 的部分。四层的 transformer block 的个数依次为（4，6，6，8），refinement部分的transformer block数量为4 。
通过 skip connection 将编码器和解码器的输出拼接，拼接时得到 C 上的维度会翻倍，然后使用 1X1 的总面积来降维。值得注意的是，最后一层（画红圈的部分），没有使用 1X1 的卷积处理。
4、Method
4.1 MDTA
与一般 Transformer 不同的是，论文在自注意力模板进行 token 计算时不是常见的 patch-wise，而是 pixel-wise。对于输入的特征首先利用 1X1 的 point-wise卷积升维，然后用 3X3 的 depth-conv 分组卷积将特征分为三块。接下来就是很常规的 self-attention 计算了。
值得注意的是，因为 W ∗ H 的维度远高于 C ，因此相似性的计算是在 C这个维度上进行的。这里也使用了 multi-head 来降低计算量来分组学习。
MDTA 模块用来替代传统 Transformer 中的多头自注意力机制，主要设计如下:
1. 多头卷积注意力机制:
- 在 MDTA 中,不是使用全连接层计算注意力权重,而是使用多组 1x1 卷积核来实现注意力计算。
- 这样可以有效降低注意力机制的计算复杂度,从而使其更适用于高分辨率图像。2. 转置卷积注意力:
- 在注意力机制中,原始 Transformer 是使用 softmax 函数将注意力权重归一化。
- 而 MDTA 则采用了"转置卷积"的方式来生成注意力权重,这种方式可以更好地保留局部空间信息。3. 多尺度融合:
- MDTA 模块使用了多组不同感受野大小的卷积核,可以捕获不同尺度的像素依赖关系。
- 这些多尺度的注意力特征被拼接在一起,形成最终的注意力输出。总的来说,MDTA 模块通过使用卷积注意力、转置卷积以及多尺度融合等方法,在大幅降低计算复杂度的同时,也能有效地建模图像中的长距离像素依赖关系。这使得 Restormer 可以应用于高分辨率图像修复任务。
## Multi-DConv Head Transposed Self-Attention (MDTA)
class Attention(nn.Module):
def __init__(self, dim, num_heads, bias):
super(Attention, self).__init__()
self.num_heads = num_heads # 注意力头的个数
self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1)) # 可学习系数
# 1*1 升维
self.qkv = nn.Conv2d(dim, dim*3, kernel_size=1, bias=bias)
# 3*3 分组卷积
self.qkv_dwconv = nn.Conv2d(dim*3, dim*3, kernel_size=3, stride=1, padding=1, groups=dim*3, bias=bias)
# 1*1 卷积
self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
def forward(self, x):
b,c,h,w = x.shape # 输入的结构 batch 数，通道数和高宽
qkv = self.qkv_dwconv(self.qkv(x))
q,k,v = qkv.chunk(3, dim=1) # 第 1 个维度方向切分成 3 块
# 改变 q, k, v 的结构为 b head c (h w)，将每个二维 plane 展平
q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
q = torch.nn.functional.normalize(q, dim=-1) # C 维度标准化，这里的 C 与通道维度略有不同
k = torch.nn.functional.normalize(k, dim=-1)
attn = (q @ k.transpose(-2, -1)) * self.temperature
attn = attn.softmax(dim=-1)
out = (attn @ v) # 注意力图(严格来说不算图)
# 将展平后的注意力图恢复
out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
# 真正的注意力图
out = self.project_out(out)
return out
python
4.2 GDFN
这里作者提出了一个新的模块来取代 Transformer block 里的 FFN 部分，是一个双路的门控网络，两个分支都是先用 1X1 的 point conv 处理，用于扩展特征通道（膨胀比为 γ = 2.66），然后两个分支都使用 3X3 的 depth conv 提取特征。下面分支使用GELU激活函数得到门控。最后用 1X1 的 point conv卷积降维输出。
参数 γ = 2.66，不清楚为什么设置为这个值，实验部分也没有分析。
有个知乎文章 里解释说：GELU可以看作 dropout的思想和relu的结合，从公式上来看，GELU对于输入乘以一个0,1组成的mask，而该mask的生成则是依靠伯努利分布的随机输入。早期的深度学习库里没有 GELU 函数，因此在较早的keras和torch的transformer实现的代码里都是用relu+dropout而没有使用gelu来引入一定的随机性，可以让神经网络对大数据的训练更鲁棒。
## Gated-Dconv Feed-Forward Network (GDFN)
class FeedForward(nn.Module):
def __init__(self, dim, ffn_expansion_factor, bias):
super(FeedForward, self).__init__()
# 隐藏层特征维度等于输入维度乘以扩张因子
hidden_features = int(dim*ffn_expansion_factor)
# 1*1 升维
self.project_in = nn.Conv2d(dim, hidden_features*2, kernel_size=1, bias=bias)
# 3*3 分组卷积
self.dwconv = nn.Conv2d(hidden_features*2, hidden_features*2, kernel_size=3, stride=1, padding=1, groups=hidden_features*2, bias=bias)
# 1*1 降维
self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
def forward(self, x):
x = self.project_in(x)
x1, x2 = self.dwconv(x).chunk(2, dim=1) # 第 1 个维度方向切分成 2 块
x = F.gelu(x1) * x2 # gelu 相当于 relu+dropout
x = self.project_out(x)
return x
python
在 GDFN 中存在一个扩展因子 γ = 2.66 \gamma=2.66γ=2.66 ，论文对其的解释很简单，就是为了
4.3 TransformerBlock
## 就是标准的 Transformer 架构
class TransformerBlock(nn.Module):
def __init__(self, dim, num_heads, ffn_expansion_factor, bias, LayerNorm_type):
super(TransformerBlock, self).__init__()
self.norm1 = LayerNorm(dim, LayerNorm_type) # 层标准化
self.attn = Attention(dim, num_heads, bias) # 自注意力
self.norm2 = LayerNorm(dim, LayerNorm_type) # 层表转化
self.ffn = FeedForward(dim, ffn_expansion_factor, bias) # FFN
def forward(self, x):
x = x + self.attn(self.norm1(x)) # 残差
x = x + self.ffn(self.norm2(x)) # 残差
return x
python
4.4 Progressive Learning
基于 CNN 的恢复模型通常在固定大小的图像块上进行训练。然而，在小裁剪补丁上训练 Transformer 模型可能不会对全局图像统计数据进行编码，从而在测试时在全分辨率图像上提供次优性能。为此，我们进行渐进式学习，其中网络在早期训练时期在较小的图像块上进行训练，在后期训练时期在逐渐增大的图像块上进行训练。
通过渐进式学习在混合大小的补丁上训练的模型在测试时表现出增强的性能，其中图像可以具有不同的分辨率（图像恢复中的常见情况）。渐进式学习策略的行为方式与课程学习过程类似，其中网络从更简单的任务开始，逐渐转向学习更复杂的任务（需要保留精细图像结构/纹理）。由于对大补丁的训练需要更长的时间，因此我们随着补丁大小的增加而减少批量大小，以保持每个优化步骤与固定补丁训练的时间相似。
训练时应用了 Progressive Learning （训练的图像逐渐增大），这是因为在裁剪的小块图像上训练时，难以获足够的全局信息，从而在测试大尺寸图像时性能不够好。因此，在训练时将 patch 逐渐从128 增大到 384 ，实验中也有地方分析了这样做确实能够涨点。
4.5 测试实例
model = Restormer()
print(model) # 打印网络结构
x = torch.randn((1, 3, 64, 64)) #随机生成输入图像
x = model(x) # 送入网络
print(x.shape) # 打印网络输入的图像结构
python
5、Experiments
实验部分没什么好说的，就是比较基础的主观+客观+消融实验，指标只用了PSNR和SSIM，没有比较网络的参数量，不清楚性能的提升是来自于网络参数的增加还是改进的两个模块。
下面标题列一下具体做的实验内容
5.1 Image deraining Results
5.2 Single-image Motion Deblurring Results
5.3 Defocus Deblurring Results
5.4 Image Denoising Results
5.5 Ablation Studies
6、Summary
我们为 Transformer 模块的核心组件引入了关键设计，以改进特征聚合和转换。具体来说，我们的 Multi-Dconv 头部转置注意力（MDTA）模块通过跨通道而不是空间维度应用自注意力来隐式模拟全局上下文，因此具有线性复杂度而不是二次复杂度。此外，所提出的门控 DConv 前馈网络（GDFN）引入了门控机制来执行受控特征转换。为了将 CNN 的优势融入到 Transformer 模型中，MDTA 和 GDFN 模块都包含用于编码空间局部上下文的深度卷积。
- 实验部分大多在介绍方法性能非常好，为什么性能会好，分析的内容相对少。
- 实验部分只比较了PSNR和SSIM，没有比较网络的参数量，无法判断性能的提升是因为网络参数的增加还是改进的两个模块。
- 随时补充嘻嘻