图像恢复 Restormer: 彻底理解论文和源代码 (注释详尽)

文章目录

1. Restormer 论文

在这里插入图片描述
主要工作：
[1] MDTA (Multi-Dconv Head Transposed Attention), 聚合局部和非局部的像素交互，可以有效处理高分辨率图像。
[2] GDFN (Gated-Dconv Feed-Forward Network), 控制特征转换，抑制小信息量的特征，仅允许有用的信息进入下一次网络。
论文：https://arxiv.org/pdf/2111.09881.pdf
源代码：
[1] https://github.com/swz30/Restormer
[2] https://download.csdn.net/download/Wenyuanbo/83592489
网络细节注释和自定义训练测试代码：https://download.csdn.net/download/Wenyuanbo/83617599

2. Restormer 网络结构

2.1 整体框架

在这里插入图片描述
论文的主要创新点是将经典 Transformer 中的 MSA 和 FFN 都进行改进，并且采用 Encoder-Decoder 架构，涉及的上采样操作使用 nn.PixelShuffle() 来实现，涉及的下采样操作使用 nn.PixelUnshuffle() 来实现，整体文章的脉络非常清晰。

2.2 MDTA

与一般 Transformer 不同的是，论文在自注意力模板进行 token 计算时不是常见的 patch-wise，而是 pixel-wise。首先利用 11 卷积升维，再使用 33 分组卷积将特征分为三块，最后进行经典的自注意力计算。
在这里插入图片描述

2.3 GDFN

论文提出双路门控网络来取代 FFN，分别进行 11 升维，再利用 33 分组卷积提取特征，随后使用 GELU 激活函数门控，最后 1*1 卷积降维输出。
在这里插入图片描述

3. 主要代码理解

3.1 MDTA

## 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

3.2 GDFN

## 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

3.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

3.4 一个测试实例

model = Restormer()
print(model)  # 打印网络结构

x = torch.randn((1, 3, 64, 64))  #随机生成输入图像
x = model(x)  # 送入网络
print(x.shape) # 打印网络输入的图像结构

参考文献

[1] Zamir S W, Arora A, Khan S, et al. Restormer: Efficient Transformer for High-Resolution Image Restoration[J]. arXiv preprint arXiv:2111.09881, 2021.
[2] 中国海洋大学AI前沿理论组. 【ARXIV2111】Restormer: Efficient Transformer for High-Resolution Image Restoration.

结语与思考

实验证明 Restormer 在图像去雨、图像去模糊、图像去噪等任务上都取得了相当不错的表现，但是论文没有与别的算法进行参数和效率的对比，据我所知 MPRNet 的参数量为 3.64 M，而 Restormer 是 25.3 M，如果靠堆参数量和烧钱使结果 SOTA 那我等小组就只能望而却步了。
在 GDFN 中存在一个扩展因子 $\gamma=2.66$ ，论文对其的解释很简单，就是为了使网络参数和计算负担与一般 FFN 一致。
完整的注释和自定义训练测试代码请移步：https://download.csdn.net/download/Wenyuanbo/83617599