SwinIR: Image Restoration Using Swin Transformer

• ICCV2021 workshop
• https://github.com/JingyunLiang/SwinIR

1. 问题引入

• 将swim transformer使用到图像恢复任务当中，因为卷积存在不能建模长距离依赖以及使用相同的卷积核来恢复不同的图像区域；
• 并不是首个将transformer引入图像恢复中的方法，但是之前的方法将图片分为固定大小的patch，在patch的边缘会产生伪影

2. 方法

• 模型分为三个部分：shallow feature extraction, deep feature extraction和high-quality image reconstruction，其中深层特征提取器包含多个residual Swin Transformer blocks (RSTB)
• 特征提取模块对所有的任务共享，但是重建模块为任务单独设计；
  
• Shallow and deep feature extraction：浅层特征提取，输入$I_{LQ}\in {\mathbb{R}}^{H\times W\times C_{in}}$，经过一个$3\times 3$的卷积层$H_{SF}(\cdot )$得到输出$F_0\in {\mathbb{R}}^{H\times W\times C}$，之后进行深层特征提取$F_{DF}\in {\mathbb{R}}^{H\times W\times C},F_{DF}=H_{DF}(F_0)$，其中$H_{DF}(\cdot )$包含$K$个residual Swin Transformer blocks (RSTB) and a 3 × 3 convolutional layer
• 图像重建层：以图像超分任务为例$I_{RHQ}=H_{REC}(F_0+F_{DF})$，使用sub-pixel convolution layer来进行上采样，假如不需要上采样的任务例如图像去噪和JPEG compression artifact reduction，就使用一层卷积实现重建层，模型学习的目标是LR和HR的residual，所以$I_{RHQ}=H_{SwinIR}(I_{LQ})+I_{LQ}$；
• 超分任务损失函数使用的是L1 pixel loss，$L=||I_{RHQ}-I_{HQ}|{|}_1$，其他任务使用的是Charbonnier loss；
• Residual Swin Transformer Block：由L个STL和一个卷积层组成；
• Swin Transformer layer：对于输入$H\times W\times C$，首先将输入reshape成$\frac{HW}{M^2}\times M^2\times C$，总共有$\frac{HW}{M^2}$个相互之间不重叠的window，之后再每个window上进行local attention，也就是标准的attention；对于当个window $X\in {\mathbb{R}}^{M^2\times C}$，$Q=X{P}_Q,K=X{P}_K,V=X{P}_V$，其中$P_Q,P_K,P_V$是不同window之间共享的参数，$Q,K,V\in {\mathbb{R}}^{M^2\times d}$，之后$Attention(Q,K,V)=SoftMax(Q{K}^T/\sqrt{d}+B)V$，其中$B$是可学习的相对位置编码，此处是多头注意力机制，接下来是MLP层(包含两层全连接和GELU)，整体$X=MSA(LN(X))+X,X=MLP(LN(X))+X$，为了保证cross-window connections，铍铜的window partition和shifted window partition是交替使用的；