目录
5.2 Single-image Motion Deblurring Results
5.3 Defocus Deblurring Results
这个论文的主要想法是将 Transformer 模型应用到图像修复中,不过和一般的VIT模型不同的是,没有使用 patch 级的特征进行 token 建模,可以理解为每个像素就是一个 token。
Transformer block 主要包括两部分:一部分是self-attention 的计算,另一部分是 FFN 部分。作者也很自然的想到了改进这两个部分,提出了两个模块,分别是:
Multi-Dconv Head Transposed Attention,主要改进 self-attention 部分
Gated-Dconv Feed-Forward Network,主要改进FFN部分
1、Motivation
卷积神经网络(CNN):在从大规模数据中学习可概括的图像先验方面表现良好,这些模型已被广泛应用于图像恢复和相关任务。
Transformers:在自然语言和高级视觉任务上表现出显着的性能提升。虽然Transformer模型减轻了CNN的缺点(有限的接收场和对输入内容的不适应性),其计算复杂度随着空间分辨率二次增长,因此不可用于涉及高分辨率图像的大多数图像恢复任务。
在这项工作中,我们提出了一个有效的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 来降低计算量来分组学习。
## 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
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
在 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
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) # 打印网络输入的图像结构
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,没有比较网络的参数量,无法判断性能的提升是因为网络参数的增加还是改进的两个模块。
- 随时补充~嘻嘻~