资讯专栏INFORMATION COLUMN

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

RiverLi / 4577人阅读

摘要:浅层特征提取浅层特征提取只使用一层卷积进行提取。图像重建模块图像重建模块其实就是卷积上采样的组合,在这块论文提出种结构。主要代码理解关于中涉及的部分代码非常简单,就不在此多带带列出,这里主要注释一下其中有关的实现代码。将维图像转变成维。

SwinIR 论文


主要工作:将 Swin Transformer 在图像恢复中应用,降低参数量的同时取得很好的效果。
论文地址https://arxiv.org/pdf/2108.10257.pdf
源代码https://github.com/JingyunLiang/SwinIR

SWinIR 网络结构

整体框架

SwinIR 的网络结构主要分为 3 个部分,分别是浅层特征提取模块,深层特征提取模块和高质量图像重建模块。其中前后两个模块都是基于 CNN 的,中间模块则主要使用 SwInTransformer。

浅层特征提取

浅层特征提取只使用一层卷积进行提取。(:这是一个容易改进的地方)

深层特征提取

深层特征提取模块由若干个残差 SwInTransformer 块 (RSTB) 和卷积块构成,具体结构如下图。

(1) 首先将来自浅层特征提取模块的特征图分割成多个不重叠的 patch embeddings;
(2) 再通过多个串联的残差 SWin Transformer 块 (RSTB);
(3) 将多个不重叠的 patch embeddings 重新组合成与输入特征图分辨率一样;
(4) 最后通过一个卷积层 (1 层或3 层卷积) 输出;
(5) 在每个 RSTB 中都引入残差连接。

残差 SwInTransformer 块 (RSTB) 中的 STL 就是 SwIn Transformer Layer 的意思,具体结构如下图。

(1) 首先通过一个归一化层 LayerNorm;
(2) 再通过多头自注意力 (Multi-head Self Attention) 模块;
(3) 在多头自注意力结尾引入残差;
(4) 再通过一个归一化层 LayerNorm;
(5) 最后通过一个多层感知机 MLP;
(6) 结尾同样引入残差。

图像重建模块

图像重建模块其实就是卷积+上采样的组合,在这块论文提出 4 种结构。(:这是一个容易改进的地方)

(1) 经典超分 (卷积 + pixelshuffle 上采样 + 卷积);
(2) 轻量超分 (卷积 + pixelshuffle 上采样);
(3) 真实图像超分 (卷积 + 卷积插值上采样 + 卷积插值上采样 + 卷积);
(4) 像去噪和 JPEG 压缩去伪影 (卷积 + 引入残差)。

主要代码理解

关于 SwinIR 中涉及 CNN 的部分代码非常简单,就不在此多带带列出,这里主要注释一下其中有关 Swin Transformer 的实现代码。另外针对不构成主要网络结构的部分代码进行了删减,完整代码请移步:
(1) GitHub 链接:https://github.com/JingyunLiang/SwinIR
(2) CSDN 链接:https://download.csdn.net/download/Wenyuanbo/40284900
(3) 详尽注释代码:https://download.csdn.net/download/Wenyuanbo/40284085

SwinIR

SwinIR 主要由浅层特征提取,深层特征提取和高质量图像重建模块组成,具体原理如前所说,直接欣赏代码吧。

# SWinIRclass SwinIR(nn.Module):    r""" SwinIR        基于 Swin Transformer 的图像恢复网络.    输入:        img_size (int | tuple(int)): 输入图像的大小,默认为 64*64.        patch_size (int | tuple(int)): patch 的大小,默认为 1.        in_chans (int): 输入图像的通道数,默认为 3.        embed_dim (int): Patch embedding 的维度,默认为 96.        depths (tuple(int)): Swin Transformer 层的深度.        num_heads (tuple(int)): 在不同层注意力头的个数.        window_size (int): 窗口大小,默认为 7.        mlp_ratio (float): MLP隐藏层特征图通道与嵌入层特征图通道的比,默认为 4.        qkv_bias (bool): 给 query, key, value 添加可学习的偏置,默认为 True.        qk_scale (float): 重写默认的缩放因子,默认为 None.        drop_rate (float): 随机丢弃神经元,丢弃率默认为 0.        attn_drop_rate (float): 注意力权重的丢弃率,默认为 0.        drop_path_rate (float): 深度随机丢弃率,默认为 0.1.        norm_layer (nn.Module): 归一化操作,默认为 nn.LayerNorm.        ape (bool): patch embedding 添加绝对位置 embedding,默认为 False.        patch_norm (bool): 在 patch embedding 后添加归一化操作,默认为 True.        use_checkpoint (bool): 是否使用 checkpointing 来节省显存,默认为 False.        upscale: 放大因子, 2/3/4/8 适合图像超分, 1 适合图像去噪和 JPEG 压缩去伪影        img_range: 灰度值范围, 1 或者 255.        upsampler: 图像重建方法的选择模块,可选择 pixelshuffle, pixelshuffledirect, nearest+conv 或 None.        resi_connection: 残差连接之前的卷积块, 可选择 1conv 或 3conv.    """    def __init__(self, img_size=64, patch_size=1, in_chans=3,                 embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,                 use_checkpoint=False, upscale=2, img_range=1., upsampler="", resi_connection="1conv",                 **kwargs):        super(SwinIR, self).__init__()        num_in_ch = in_chans  # 输入图片通道数        num_out_ch = in_chans  # 输出图片通道数        num_feat = 64  # 特征图通道数        self.img_range = img_range  # 灰度值范围:[0, 1] or [0, 255]        if in_chans == 3:  # 如果输入是RGB图像            rgb_mean = (0.4488, 0.4371, 0.4040)  # 数据集RGB均值            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)  # 转为[1, 3, 1, 1]的张量        else:  # 否则灰度图            self.mean = torch.zeros(1, 1, 1, 1)  # 构造[1, 1, 1, 1]的张量        self.upscale = upscale  # 图像放大倍数,超分(2/3/4/8),去噪(1)        self.upsampler = upsampler # 上采样方法        self.window_size = window_size  # 注意力窗口的大小        #######################################################################################        ################################### 1, 浅层特征提取 ###################################        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)  # 输入卷积层        ##########################################################################################        ################################### 2, 深层特征提取 ######################################        self.num_layers = len(depths)  # Swin Transformer 层的个数        self.embed_dim = embed_dim  # 嵌入层特征图的通道数        self.ape = ape  # patch embedding 添加绝对位置 embedding,默认为 False.        self.patch_norm = patch_norm  # 在 patch embedding 后添加归一化操作,默认为 True.        self.num_features = embed_dim  # 特征图的通道数        self.mlp_ratio = mlp_ratio  # MLP隐藏层特征图通道与嵌入层特征图通道的比        # 将图像分割成多个不重叠的patch        self.patch_embed = PatchEmbed(            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,            norm_layer=norm_layer if self.patch_norm else None)        num_patches = self.patch_embed.num_patches  # 分割得到patch的个数        patches_resolution = self.patch_embed.patches_resolution  # 分割得到patch的分辨率        self.patches_resolution = patches_resolution        # 将多个不重叠的patch合并成图像        self.patch_unembed = PatchUnEmbed(            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,            norm_layer=norm_layer if self.patch_norm else None)        # 绝对位置嵌入        if self.ape:            # 结构为 [1,patch个数, 嵌入层特征图的通道数] 的参数            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))            trunc_normal_(self.absolute_pos_embed, std=.02)  # 截断正态分布,限制标准差为0.02        self.pos_drop = nn.Dropout(p=drop_rate)  # 以drop_rate为丢弃率随机丢弃神经元,默认不丢弃        # 随机深度衰减规律,默认为 [0, 0.1] 进行24等分后的列表        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]        # Residual Swin Transformer blocks (RSTB)        # 残差 Swin Transformer 块 (RSTB)        self.layers = nn.ModuleList()  # 创建一个ModuleList实例对象,也就是多个 RSTB        for i_layer in range(self.num_layers):  # 循环 Swin Transformer 层的个数次            # 实例化 RSTB            layer = RSTB(dim=embed_dim,                         input_resolution=(patches_resolution[0],                                           patches_resolution[1]),                         depth=depths[i_layer],                         num_heads=num_heads[i_layer],                         window_size=window_size,                         mlp_ratio=self.mlp_ratio,                         qkv_bias=qkv_bias, qk_scale=qk_scale,                         drop=drop_rate, attn_drop=attn_drop_rate,                         drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],  # no impact on SR results                         norm_layer=norm_layer,                         downsample=None,                         use_checkpoint=use_checkpoint,                         img_size=img_size,                         patch_size=patch_size,                         resi_connection=resi_connection                         )            self.layers.append(layer)  # 将 RSTB 对象插入 ModuleList 中        self.norm = norm_layer(self.num_features)  # 归一化操作,默认 LayerNorm        # 在深层特征提取网络中加入卷积块,保持特征图通道数不变        if resi_connection == "1conv":  # 1层卷积            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)        elif resi_connection == "3conv":  # 3层卷积            # 为了减少参数使用和节约显存,采用瓶颈结构            self.conv_after_body = nn.Sequential(nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),  # 降维                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),                                                 nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),                                                 nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1))  # 升维        # 高质量图像重建模块        if self.upsampler == "pixelshuffle":  # pixelshuffle 上采样            # 适合经典超分            self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),                                                      nn.LeakyReLU(inplace=True))            self.upsample = Upsample(upscale, num_feat)  # 上采样            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)  # 输出卷积层        elif self.upsampler == "pixelshuffledirect":  # 一步是实现既上采样也降维            # 适合轻量级充分,可以减少参数量(一步是实现既上采样也降维)            self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch,                                            (patches_resolution[0], patches_resolution[1]))        elif self.upsampler == "nearest+conv":  # 最近邻插值上采样            # 适合真实图像超分            assert self.upscale == 4, "only support x4 now."  # 声明目前仅支持4倍超分重建            # 上采样之前的卷积层            self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),                                                      nn.LeakyReLU(inplace=True))            # 第一次上采样卷积(直接对输入做最近邻插值变为2倍图像)                                                      self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)            # 第二次上采样卷积(直接对输入做最近邻插值变为2倍图像)            self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)            self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)  # 对上采样完成的图像再做卷积            self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)  # 激活层            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)  # 输出卷积层        else:            # 适合图像去噪和 JPEG 压缩去伪影            self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)        self.apply(self._init_weights)  # 初始化网络参数    # 初始化网络参数    def _init_weights(self, m):        if isinstance(m, nn.Linear):  # 判断是否为线性 Linear 层            trunc_normal_(m.weight, std=.02)  # 截断正态分布,限制标准差为 0.02            if m.bias is not None:  # 如果设置了偏置                nn.init.constant_(m.bias, 0)  # 初始化偏置为 0        elif isinstance(m, nn.LayerNorm):  # 判断是否为归一化 LayerNorm 层            nn.init.constant_(m.bias, 0)  # 初始化偏置为 0            nn.init.constant_(m.weight, 1.0)  # 初始化权重系数为 1        # 检查图片(准确说是张量)的大小    def check_image_size(self, x):        _, _, h, w = x.size()  # 张量 x 的高和宽        # h 维度要填充的个数        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size        # w 维度要填充的个数        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size        # 右填充 mod_pad_w 个值,下填充 mod_pad_h 个值,模式为反射(可以理解为以 x 的维度末尾为轴对折)        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect")        return x    # 深层特征提取网络的前向传播    def forward_features(self, x):        x_size = (x.shape[2], x.shape[3])  # 张量 x 的高和宽        x = self.patch_embed(x)  # 分割 x 为多个不重叠的 patch embeddings        if self.ape:  # 绝对位置 embedding            x = x + self.absolute_pos_embed  # x 加上对应的绝对位置 embedding        x = self.pos_drop(x)  # 随机将x中的部分元素置 0        for layer in self.layers:            x = layer(x, x_size)  # x 通过多个串联的 RSTB        x = self.norm(x)  # 对 RSTB 的输出进行归一化        x = self.patch_unembed(x, x_size)  # 将多个不重叠的 patch 合并成图像        return x        # SWinIR 的前向传播    def forward(self, x):        H, W = x.shape[2:]  # 输入图片的高和宽        x = self.check_image_size(x)  # 检查图片的大小,使高宽满足 window_size 的整数倍                self.mean = self.mean.type_as(x) # RGB 均值的类型同 x 一致        x = (x - self.mean) * self.img_range  # x 减去 RGB 均值再乘以输入的最大灰度值        if self.upsampler == "pixelshuffle":  # pixelshuffle 上采样方法            # 适合经典超分            x = self.conv_first(x)  # 输入卷积层            x = self.conv_after_body(self.forward_features(x)) + x  # 深度特征提取网络,引入残差            x = self.conv_before_upsample(x)  # 上采样前进行卷积            x = self.conv_last(self.upsample(x))  # 上采样后再通过输出卷积层        elif self.upsampler == "pixelshuffledirect":  # 一步是实现既上采样也降维            # 适合轻量级超分            x = self.conv_first(x)  # 输入卷积层            x = self.conv_after_body(self.forward_features(x)) + x  # 深度特征提取网络,引入残差            x = self.upsample(x)  # 上采样并降维后输出        elif self.upsampler == "nearest+conv":  # 最近邻插值上采样方法            # 适合真实图像超分,只适合 4 倍超分            x = self.conv_first(x)  # 输入卷积层            x = self.conv_after_body(self.forward_features(x)) + x  # 深度特征提取网络,引入残差            x = self.conv_before_upsample(x)  # 上采样前进行卷积            # 第一次上采样 2 倍            x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest")))            # 第二次上采样 2 倍            x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest")))            x = self.conv_last(self.lrelu(self.conv_hr(x))) #  输出卷积层        else:            # 适合图像去噪和 JPEG 压缩去伪影            x_first = self.conv_first(x)  # 输入卷积层            res = self.conv_after_body(self.forward_features(x_first)) + x_first  # 深度特征提取网络,引入残差            x = x + self.conv_last(res) #  输出卷积层,引入残差        x = x / self.img_range + self.mean  # 最后的 x 除以灰度值范围再加上 RGB 均值        return x[:, :, :H*self.upscale, :W*self.upscale]  # 返回输出 x

MLP

多层感知机 MLP 是 transformer 比较基础的部分,具体原理也很简单。

# 多层感知机class Mlp(nn.Module):    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):        super().__init__()        out_features = out_features or in_features  # 输入特征的维度        hidden_features = hidden_features or in_features  # 隐藏特征维度        self.fc1 = nn.Linear(in_features, hidden_features)  # 线性层        self.act = act_layer()  # 激活函数        self.fc2 = nn.Linear(hidden_features, out_features)  # 线性层        self.drop = nn.Dropout(drop)  # 随机丢弃神经元,丢弃率默认为 0        # 定义前向传播    def forward(self, x):        x = self.fc1(x)        x = self.act(x)        x = self.drop(x)        x = self.fc2(x)        x = self.drop(x)        return x

Patch Embedding

主要的操作就是将原始 2 维图像 (特征图的一个 plane 或者说一个 channel) 转变为 1 维的 patch embeddings,通过 Swin Transformer 学习处理之后再重新组合成与原来特征图结构一致的新特征图。

(1) 将 2 维图像转变成 1 维 patch embeddings。

# 图像转成 Patch Embeddingsclass PatchEmbed(nn.Module):    
                 
               
              

文章版权归作者所有,未经允许请勿转载,若此文章存在违规行为,您可以联系管理员删除。

转载请注明本文地址:https://www.ucloud.cn/yun/123162.html

相关文章

  • 【超分辨率】略读21年超分SOTA新算法——SwinIR

    摘要:要点目前超分的算法是基于的,几乎没有用的研究。要点网络结构如下任务训练时损失函数采用损失,对于图像去噪和压缩伪影重建任务上,使用。要点源码实现其实就是全连接网络,用了损失函数的全称是与相比,像,和这样的激活可以使神经网络更快更好地收敛。 ...

    Alliot 评论0 收藏0
  • 论文解读 用于弱监督表面缺陷分割的缺陷注意模板循环对抗网络 (Defect attention te

    摘要:本文提出了一种弱监督缺陷分割方法,该方法基于由图像级注释训练的改进的循环一致生成对抗网络生成的动态模板。一些研究人员将弱监督学习应用于表面缺陷检测。生成器生成器的主要目的是生成对应于缺陷图像的无缺陷模板。 ...

    jkyin 评论0 收藏0
  • 躺平吧,平铺的窗口「GitHub 热点速览 v.21.47」

    摘要:作者小鱼干用系统经常会遇到的一个问题便是多开窗口如何快速找寻的问题,本周特推项目便是来解决这个问题的。以下内容摘录自微博的及热帖简称热帖,选项标准新发布实用有趣,根据项目时间分类,发布时间不超过的项目会标注,无该标志则说明项目超过半月。 作者:HelloGitHub-小鱼干 用 macOS 系统经常会遇到的一个问题便是多开窗口如何快速找寻的问题,本周特推项目 yabai 便是来解...

    galaxy_robot 评论0 收藏0
  • 何恺明团队推出Mask^X R-CNN,将实例分割扩展到3000类

    摘要:从标题上可以看出,这是一篇在实例分割问题中研究扩展分割物体类别数量的论文。试验结果表明,这个扩展可以改进基准和权重传递方法。 今年10月,何恺明的论文Mask R-CNN摘下ICCV 2017的较佳论文奖(Best Paper Award),如今,何恺明团队在Mask R-CNN的基础上更近一步,推出了(以下称Mask^X R-CNN)。这篇论文的第一作者是伯克利大学的在读博士生胡戎航(清华...

    MockingBird 评论0 收藏0
  • 异构传感器定位论文概述

    摘要:一简介最近在研究异构传感器定位方面的东西,读了一些论文。这个方向的一个基本逻辑就是,先用激光雷达等高精度设备建立一个高精度的点云地图,然后再用视觉里程计结合异构传感器定位的一些方法,实现一个累计误差可控的视觉里程计。 ...

    HelKyle 评论0 收藏0

发表评论

0条评论

最新活动
阅读需要支付1元查看
<