EDT：针对Low-Level视觉任务的一种高效Transformer基准框架

EDT：On Efficient Transformer-Based Image Pre-training for Low-Level Vision

论文地址：On Efficient Transformer-Based Image Pre-training for Low-Level Vision

代码地址：fenglinglwb/EDT: On Efficient Transformer-Based Image Pre-training for Low-Level Vision

现阶段问题

​ 预训练在high-level计算机视觉中产生了许多最先进的技术，但很少有人尝试研究预训练如何在low-level任务中。

主要贡献

1. 提出了一个用于低级视觉的高效且通用的Transformer框架：改进window attention，分别从高、宽进行切块计算注意力。

2. 是第一个对低级视觉的图像预训练进行深入研究的人，揭示了预训练如何影响模型的内部表示以及如何进行有效的预训练的见解

网络框架

Shifted Crossed Local Attention

CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows基础上进行改进。
$$\begin{array}{rl}& \mathbf{X}=[{\mathbf{X}}_1,{\mathbf{X}}_2],\mathrm{where}{\mathbf{X}}_1,{\mathbf{X}}_2\in {\mathbb{R}}^{(H\times W){\times }^{C/2}},\\ & {\mathbf{X}}_1^{{}^{\prime }}=\mathrm{H}-\mathrm{MSA}({\mathbf{X}}_1),\\ & {\mathbf{X}}_2^{{}^{\prime }}=\mathrm{V}-\mathrm{MSA}({\mathbf{X}}_2),\\ & (\mathrm{S})\mathrm{CL}-\mathrm{MSA}(\mathbf{X})=\mathrm{Proj}([{\mathbf{X}}_1^{{}^{\prime }},{\mathbf{X}}_2^{{}^{\prime }}]),\end{array}$$

    def calculate_mask(self, x_size, index):
        # calculate attention mask for SW-MSA
        if self.shift_size is None:
            return None

        H, W = x_size
        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
        h_window_size, w_window_size = self.window_size[0], self.window_size[1]
        h_shift_size, w_shift_size = self.shift_size[0], self.shift_size[1]
        if index == 1:
            h_window_size, w_window_size = self.window_size[1], self.window_size[0]
            h_shift_size, w_shift_size = self.shift_size[1], self.shift_size[0]

        h_slices = (slice(0, -h_window_size),
                    slice(-h_window_size, -h_shift_size),
                    slice(-h_shift_size, None))
        w_slices = (slice(0, -w_window_size),
                    slice(-w_window_size, -w_shift_size),
                    slice(-w_shift_size, None))
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(img_mask, self.window_size, index)  # nW, h_window_size, w_window_size, 1
        mask_windows = mask_windows.view(-1, h_window_size * w_window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

        return attn_mask

    x = self.qkv(x) #B H W 3C
    x = x.view(B, H, W, 3, C).permute(3,0,1,2,4).contiguous()#3 B H W C
    x_h = x[...,:C//2]
    x_v = x[...,C//2:]

    if self.shift_size:
        x_h = torch.roll(x_h, shifts=(-self.shift_size[0],-self.shift_size[1]), dims=(2,3))
        x_v = torch.roll(x_v, shifts=(-self.shift_size[1],-self.shift_size[0]), dims=(2,3))

    if self.input_resolution == x_size:
        attn_windows_h = self.attns[0](x_h, mask=self.attn_mask_h)
        attn_windows_v = self.attns[1](x_v, mask=self.attn_mask_v)
    else:
        mask_h = self.calculate_mask(x_size, index=0).to(x_h.device) if self.shift_size else None
        mask_v = self.calculate_mask(x_size, index=1).to(x_v.device) if self.shift_size else None
        attn_windows_h = self.attns[0](x_h, mask=mask_h)
        attn_windows_v = self.attns[1](x_v, mask=mask_v)

    if self.shift_size:
        attn_windows_h = torch.roll(attn_windows_h, shifts=(self.shift_size[0],self.shift_size[1]), dims=(1,2))
        attn_windows_v = torch.roll(attn_windows_v, shifts=(self.shift_size[1],self.shift_size[0]), dims=(1,2))

    attn_windows = torch.cat([attn_windows_h, attn_windows_v], dim=-1)
    attn_windows = self.proj(attn_windows) #B H W C

Anti-Blocking FFN

​ 分组卷积是指将输入和输出通道分为若干组，在每组内部进行卷积操作，这可以加速计算并在一定程度上提高模型的表征能力

self.dwconv = nn.Conv2d(hidden_features, hidden_features, 5, 1, 5//2, groups=hidden_features)

Convolution Block