HAT论文详解：Activating More Pixels in Image Super-Resolution Transformer

code：https://github.com/XPixelGroup/HAT
paper: https://arxiv.org/abs/2309.05239

1. 概述

本文是对Swinir的改进，目前很多图像超分Benchmark的SOTA。相对于SwinIR的改进主要有三个地方：1. 引入Channel Attention,以获得更好的全局能力；2. 提出了overlapping cross-attention模块，来进行跨window的信息交互；3. 提出一个预训练策略。

2. 引言

2.1 阐明swinir存在的问题

SwinIR在SR任务上取得了突破，然而为什么Transformer-based方法要比CNN-based方法好，却很难说清楚。一个直观的解释是Transformer方法可以受益于self-attention机制，并能够利用远距离信息。作者通过LAM分析发现，与RCAN相比，SwinIR并没有利用更大range的信息，这是反直觉的。同时可以说明SwinIR具备比CNN强的映射能力，可以利用更少的信息取得更好的效果。但是由于利用的pixel的范围有限，SwinIR可能会restore出错误的纹理。如下图所示。
尽管平均性能优于RCAN，但是有一些结果也比RCAN差
这说明Swin transformer建模局部信息的能力很强，但是探索的信息范围需要扩大
在SwinIR的特征图上发现了block artifacts，这是由于窗口划分造成的，这说明移动窗口机制并不能有效的建立跨窗口的交互。。

2.2 本文的贡献：

设计了一个Hybrid Attention,结合了channel attention， self-attention和overlapping cross-attention；
channel attention：具备很好地获取全局信息的能力
self-attention: 强大的表达能力（representative ability）
提出一个预训练策略
因为transformer不具备cnn的归纳偏置，所以需要大规模数据进行预训练，才能解锁潜力。

3. 方法介绍

HAT结构图
在这里插入图片描述
上面两张图分别是HAT和SwinIR的整体结构图，可以看出HAT延续了SwinIR的基本结构，将RSTB升级成RHAG，内部的STL也对应升级成HAB，并且在每个Block中加入了一个OCAB。下面具体来看这两处改动。

向(STL)Swin Transformer Layer中加入了Channel Attention，也就是将(S)W-MSA与CAB的结果叠加起来。

CAB的代码实现：

class ChannelAttention(nn.Module):
    """Channel attention used in RCAN.
    Args:
        num_feat (int): Channel number of intermediate features.
        squeeze_factor (int): Channel squeeze factor. Default: 16.
    """

    def __init__(self, num_feat, squeeze_factor=16):
        super(ChannelAttention, self).__init__()
        self.attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),
            nn.ReLU(inplace=True),
            nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0),
            nn.Sigmoid())

    def forward(self, x):
        y = self.attention(x)
        return x * y

class CAB(nn.Module):

    def __init__(self, num_feat, compress_ratio=3, squeeze_factor=30):
        super(CAB, self).__init__()

        self.cab = nn.Sequential(
            nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1),
            nn.GELU(),
            nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1),
            ChannelAttention(num_feat, squeeze_factor)
            )

    def forward(self, x):
        return self.cab(x)

在每一个RHAG的最后引入一个Overlapping Cross-Attention Block (OCAB)，直接建立跨窗口的连接，同时增强窗口自注意力的表达能力。实现方式仍是基于W-MSA，只是在窗口划分时，Q的窗口是正常的无overlap的窗口，窗口大小为M * M，而K和V的窗口大小是M0 * M0， M0 =（1+gamma) * M, gamma是用于控制重叠大小的参数。虽然窗口的大小不一样，但是窗口的数量是相同的，一一对应的。
```
  Q shape: (nums_of_windows, M*M, emb_dims)
  
  K shape: (nums_of_windows, M0*M0, emb_dims)
  
  V shape: (nums_of_windows, M0*M0, emb_dims)
  
  QK.T shape: (nums_of_windows, M*M, M0*M0)
  
  因此得到的结果仍是 (nums_of_windows, M*M, emb_dims)，但其过程中获取了跨窗口的信息，因为OCA的key和value是从更大的区域中计算得到的，因此更多有用的信息将被query查询到；
```
- 预训练
  
  使用Imagenet进行X4预训练，再在DF2K上进行finetune, 发现很有效，预训练的效果取决于数据的量级和多样性；同时，作者指出充分的iteration和合适的小学习率对于预训练来说非常重要；