Restormer: Efficient Transformer for High-Resolution Image Restoration 论文阅读

date

Mar 13, 2023

Last edited time

Mar 27, 2023 08:26 AM

status

Published

slug

Restormer论文阅读

Abstract

在这项工作中，作者提出了一个高效的 Transformer 模型——Restoration Transformer（Restormer），通过在构建模块（multi-head attention and feed-forward network）中进行几个关键设计，使其能够捕获长距离像素交互的同时，仍然适用于大分辨率图像

卷积神经网络（CNN）在从大规模数据中学习广义图像先验知识方面表现良好，这些模型已被广泛应用于图像恢复和相关任务
CNN 中的基本操作是 “卷积”，它提供了 local connectivity 局部连通性和 translation equivariance 翻译等价性，这些特性为 CNN 带来了效率和通用性

两个主要问题

感受野有限，因此无法对长距离像素相关性进行建模
卷积滤波器在推理时具有静态权重，因此不能灵活地适应输入内容

为了解决上述缺点，一种更强大、更动态的替代方法是自注意（SA）机制，该机制通过所有其他位置的加权和计算给定像素处的响应: Transformers：Self-attention

Transformer 模型缓解了 CNN 的缺点（有限的感受野和对输入内容的不适应性）

SA 在捕捉远距离像素交互方面非常有效，但其计算复杂度随着空间分辨率的增加而呈二次方增长，因此无法应用于高分辨率图像

为了减少计算量，有些方法是将图像分解为一系列 patch（局部窗口），并了解它们之间的相互关系。这些模型的显著特点是具有强大的能力，能够学习到图像 patch 之间的长距离依赖关系以及对给定输入内容的适应性

此外，还可以通过以下方法来减少计算量：

在每个像素周围大小为 8×8 的小空间窗口上应用 SA

将输入图像分割为大小为 48×48 的非重叠 patch，并独立计算每个 patch 上的 SA

但这种限制 SA 的空间范围与捕获真正的远程空间像素关系的目标自相矛盾，尤其是在高分辨率图像上。本文作者提出了一种高效的图像恢复 transformer ，它能够建模全局连通性，并且仍然适用于高分辨率图像:

引入了一个 multi-Dconv head ‘transposed’ attention (MDTA) block 来代替普通多头 SA，将 SA 应用于通道维度而不是空间维度。

feed-forward network (FN) 由两个全连接层组成，层间具有非线性。本文作者使用门控机制重新制定常规 FN 的第一个线性变换层，以改善通过网络的信息流

渐进式学习策略。网络在早期的小 patch 和大 batch 上进行训练，在后期逐渐在大 patch 和小 batch 上进行训练

Method

本文主要目标是开发一个高效的 transformer 模型，该模型可以处理用于恢复任务的高分辨率图像。为了缓解计算瓶颈，引入了 multi-Dconv head transposed attention (MDTA) 和 Gated-Dconv feed-forward network (GDFN) 的关键设计。

💡

Encoder-Decoder

输入退化图像，Restormer 首先卷积得到低级特征嵌入

浅层特征通过 4 个层对称的 encoder-decoder 来转化为深层特征

从高分辨率输入开始，编码器分层减少空间大小，同时扩展通道数
解码器采用低分辨率潜在特征作为输入，并逐层恢复高分辨率表示
对于特征上下采样，作者分别应用了 pixel-unshuffle and pixel-shuffle 操作

💡

PixelShuffle : Rearranges elements in a tensor of shape to a tensor of shape , where is an upscale factor.

[PixelUnshuffle] : Reverses the PixelShuffle operation by rearranging elements in a tensor of shape to a tensor of shape , where is a downscale factor.

通过跳跃连接来拼接编码器和解码器输出，拼接后使用 1x1 卷积来减少通道数（减半）

Multi-Dconv Head Transposed Attention(MDTA)

Transformer中的主要计算开销来自于self-attention，其计算复杂度与输入图片的尺寸的平方成正比，对于尺寸为的输入图像，其计算复杂度为，所以在需要输入高分辨率图片的应用中很难应用SA模块。

MDTA模块拥有线性复杂度。关键是跨通道而不是空间维度应用SA，这儿应该是指矩阵乘法的运算结果维度为 , 即计算跨通道的交叉协方差，以生成一个对全局上下文进行隐式编码的注意图。作为MDTA中的另一个重要组件，我们在计算特征协方差以生成全局注意图之前，引入深度卷积来强调局部上下文。

💡

MDTA

对于输入，归一化得到，然后生成 query(Q), key(K), value(V)

应用 1×1 point-wise 卷积来聚合像素级跨通道上下文，然后应用 3×3 depth-wise 卷积来编码通道级空间上下文

之后对 query 和 key 做一下 reshape 操作，来 dot-product 计算生成 transposed-attention map ，而不是传统方法生成的注意力图

最终的输出为

Gated-Dconv Feed-Forward Network(GDFN)

在前馈网络部分，我们提出了两个基本改进来改进表示学习:(1)门控机制和(2)深度卷积。我们的GDFN架构如下图所示

门控机制被制定为两个平行路径的线性变换层的元件的乘积，其中一个被GELU非线性激活。它使用两个1×1卷积，一个用于扩展特征通道（膨胀比），另一个用于将通道缩减回原始输入维度。与在MDTA中一样，我们还在GDFN中包含深度卷积来编码信息空间上相邻像素的位置。给定一个输入张量 , GDFN表示为:

式中为元素相乘，表示GELU非线性激活层，LN为层归一化。GDFN 控制各层中的通道中的信息流，从而使得每层都专注于与其他层之间互补的精细细节。既与 MDTA 相比，GDFN 更专注于使用上下文信息丰富特性。由于 GDFN 比常规的 FN 执行更多操作，所以减低的膨胀比 γ ，以便具有相似的参数和计算量

Progressive Learning

基于cnn的复原模型通常在固定大小的图像块上进行训练。然而，在小的裁剪patch上训练Transformer模型可能无法编码全局图像统计信息，从而在测试时提供对全分辨率图像的次优性能。

为此，我们进行了渐进式学习，即网络在早期的训练时间内对较小的图像块进行训练，并在后期的训练时间内逐渐对较大的图像块进行训练。通过渐进式学习对混合大小的patch进行训练的模型在测试时表现出增强的性能，图像可以是不同的分辨率(图像恢复中的常见情况)。

渐进学习策略的行为方式类似于课程学习过程，网络从一个更简单的任务开始，然后逐渐转移到学习一个更复杂的任务(其中需要保存精细的图像结构/纹理)。由于在大patch上进行训练需要花费更长的时间，所以我们随着patch大小的增加而减少批大小，以保持每个优化步骤与固定patch训练的时间相似。

Experiments and Analysis

作者在四个图像处理任务的基准数据集和实验设置上进行了实验评估: 图像去雨，单图像运动去模糊，散焦去模糊（在单图像和双像素数据上），图像去噪（在合成和真实数据上）

💡

Detail

从第一层到第四层，Transformer blocks 数量设置为 [4，6，6，8]，MDTA 中的注意头为 [1，2，4，8]，通道的数量为 [48，96，192，384]

patch size 和 batch size [(1282, 64), (1602, 40), (1922, 32), (2562, 16), (3202, 8), (3842, 8)]