详解SwinIR的论文和代码(SwinIR: Image Restoration Using Swin Transformer)

2023-11-20 16:12

本文主要是介绍详解SwinIR的论文和代码(SwinIR: Image Restoration Using Swin Transformer),希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

paper:https://arxiv.org/abs/2108.10257
code:https://github.com/JingyunLiang/SwinIR

目录

  • 1. Swin Transformer layers
    • 1.1 局部注意力
    • 1.2 移动窗口机制
    • 1.3 关键代码理解
  • 2. 整体网络结构
    • 2.1 浅层特征提取
    • 2.2 深层特征提取
    • 2.3 图像重建
  • 3.总结

SwinIR将Swin transformer1应用到low level领域的图像增强任务,结合卷积设计了网络结构,在以下三个任务上取得了很好的效果:图像超分辨率(包括classical、lightweight和real-world SR)、图像去噪(包括灰度图和彩色图像去噪)和 JPEG压缩失真去除。本文将结合代码对SwinIR进行详解。

SwinIR的网络结构并不复杂,关键部件就是Swin Transformer layers(STL)卷积层残差连接。卷积和残差连接大家都比较熟悉了,因此我首先结合代码介绍一下swin transformer层,然后自底向上的介绍SwinIR的全貌

1. Swin Transformer layers

SwinIR使用的Swin Transformer layers(STL)是在swin transformer中提出的,并未有改动。STL基于原始的多头注意力transformer层进行优化,主要的不同点在于:1. 局部注意力(local attention);2. 移动窗口机制(shifted window mechanism);

1.1 局部注意力

原始的全局注意力会将图像分成若干个patch,所有的patch之间做自注意力计算;所谓的局部注意力就是首先将图像划分成若干个window,每个window内在进行patch的划分,然后在window内部进行自注意力的计算,而不在一个window内的patch是没有交互的。也就是说,只考虑一个window内的patch,他们之间的计算和全局注意力操作是一样的。

理解局部注意力具体是怎么做的,很好的一个办法是看代码和分析tensor在不同层之间的shape整理出来。下面是我整理的tensor shape变化:

请添加图片描述
其中,b: batchsize, h: 输入高, w:输入宽, ws: 窗口大小, C: channel数, num_heads:attention的head数

1.2 移动窗口机制

由于基于窗口的多头注意力(W-MSA)没有考虑跨窗口的连接,模型建模长距离关联的能力受损。因此swin transformer提出了移动窗口多头注意力机制(SW-MSA),可在保证计算高效性的前提下,扩大感受野。

如下图所示,W-MSA的窗口大小为M*M(图中M=4),那么SW-MSA的窗口划分将向右下移动 ⌊ M / 2 ⌋ ∗ ⌊ M / 2 ⌋ \lfloor M/2 \rfloor *\lfloor M/2 \rfloor M/2M/2

请添加图片描述

但是经过位移之后,窗口数量会变多,由原来的 ⌊ h / M ⌋ ∗ ⌊ w / M ⌋ \lfloor h/M \rfloor *\lfloor w/M \rfloor h/Mw/M变成 ( ⌊ h / M ⌋ + 1 ) ∗ ( ⌊ w / M ⌋ + 1 ) (\lfloor h/M \rfloor + 1) *(\lfloor w/M \rfloor +1) (⌊h/M+1)(⌊w/M+1),而且窗口大小不一致。因此swin transformer提出了循环位移,减少窗口数量,同时可以获得相同大小的窗口进行并行计算。循环位移如下图所示。
请添加图片描述

在代码中,循环位移通过torch.roll实现,shifts为负,代表从下往上移动,从右往左移动,最上和最左循环移动到最下和最右。

# cyclic shift
if self.shift_size > 0:shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))

关于torch.roll可参考:https://blog.csdn.net/weixin_42899627/article/details/116095067
如上图所示,经过循环移位后,有三个窗口中有一些patch是本不相邻的,它们不应该做自注意力,所以swin transformer建立了mask机制来完成最终的注意力计算。

关于mask的理解可参考https://github.com/microsoft/Swin-Transformer/issues/38

1.3 关键代码理解

下面来看一下关键代码及注释,首先是WindowAttention的forward函数:

def forward(self, x, mask=None):"""Args:x: input features with shape of (num_windows*B, N, C)mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None"""B_, N, C = x.shape  # 此处的输入是经过window partition的# self.qkv(x): num_windows*B, window_size*window_size, 3*Cqkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) # 通过一个全连接层获取所有头的qkv,(3, num_windows*B, num_heads, window_size*window_size, C // num_heads)q, k, v = qkv[0], qkv[1], qkv[2] q = q * self.scaleattn = (q @ k.transpose(-2, -1)) # num_windows*B, num_heads, window_size*window_size, window_size*window_size# 可学习的相对位置biasrelative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Wwattn = attn + relative_position_bias.unsqueeze(0) # num_windows*B, num_heads, window_size*window_size, window_size*window_sizeif mask is not None:nW = mask.shape[0]# 将mask和attn相加,mask只有两种取值0和-100,因此为0时对attn无影响,为-100时,self.softmax(attn)将变为接近于0attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)attn = attn.view(-1, self.num_heads, N, N) # num_windows*B, num_heads, window_size*window_size, window_size*window_sizeattn = self.softmax(attn) # num_windows*B, num_heads, window_size*window_size, window_size*window_sizeelse:attn = self.softmax(attn)attn = self.attn_drop(attn)# v:num_windows*B, num_heads, window_size*window_size, C // num_heads# attn:num_windows*B, num_heads, window_size*window_size, window_size*window_size# attn @ v: num_windows*B, num_heads, window_size*window_size, C // num_headsx = (attn @ v).transpose(1, 2).reshape(B_, N, C) # num_windows*B, window_size*window_size, Cx = self.proj(x) # 全连接层x = self.proj_drop(x)return x

接下来是SwinTransformerBlock的forward函数

    def forward(self, x, x_size):H, W = x_sizeB, L, C = x.shape# assert L == H * W, "input feature has wrong size"shortcut = xx = self.norm1(x)x = x.view(B, H, W, C)# cyclic shiftif self.shift_size > 0:shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))else:shifted_x = x# partition windowsx_windows = window_partition(shifted_x, self.window_size)  # (num_windows*B, window_size, window_size, C)x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # num_windows*B, window_size*window_size, C# W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window sizeif self.input_resolution == x_size:attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, Celse:attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))# merge windowsattn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C# reverse cyclic shiftif self.shift_size > 0:x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))else:x = shifted_xx = x.view(B, H * W, C)# FFNx = shortcut + self.drop_path(x)x = x + self.drop_path(self.mlp(self.norm2(x)))return x

可以看到每个SwinTransformerBlock内部完成的是:
X = M S A ( L N ( X ) ) + X X = MSA(LN(X)) + X X=MSA(LN(X))+X
X = M L P ( L N ( X ) ) + X X = MLP(LN(X)) + X X=MLP(LN(X))+X
其中MSA为W-MSA和SW-MSA交替。

2. 整体网络结构

请添加图片描述
如上图所示,SwinIR包括三个modules,浅层特征提取、深层特征提取和图像重建。其中特征提取模块对所有任务都是一样的,但是图像重建对于不同的任务是不同的。

2.1 浅层特征提取

一个3×3卷积层将特征图通道转成embed_dim:(b, embed_dim, h, w)

self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)

2.2 深层特征提取

深层特征提取的基本模块则是第一节中讲解的STL和卷积层和残差连接。STL和卷积组成RSTB,RSTB和卷积组成了深层特征提取。

2.3 图像重建

以下代码可以看到对于不同的任务,图像重建模块是不同的,有的采用最邻近插值+卷积,有的采用pixelshuffle+卷积,有的直接采用卷积。


if self.upsampler == 'pixelshuffle':# for classical SRx = self.conv_first(x)x = self.conv_after_body(self.forward_features(x)) + xx = self.conv_before_upsample(x)x = self.conv_last(self.upsample(x))
elif self.upsampler == 'pixelshuffledirect':# for lightweight SRx = self.conv_first(x)x = self.conv_after_body(self.forward_features(x)) + xx = self.upsample(x)
elif self.upsampler == 'nearest+conv':# for real-world SRx = self.conv_first(x) # (b, embed_dim, h, w)x = self.conv_after_body(self.forward_features(x)) + xx = self.conv_before_upsample(x)x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))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:# for image denoising and JPEG compression artifact reductionx_first = self.conv_first(x)res = self.conv_after_body(self.forward_features(x_first)) + x_firstx = x + self.conv_last(res)

SwinIR可以很灵活配置网络的复杂度。影响W-MSA计算复杂度: 4 h w C 2 + 2 M 2 h w C 4hwC^2 + 2M^2hwC 4hwC2+2M2hwC
请添加图片描述

3.总结

  1. 结构简单,性能全面超过cnn-based的方法,适用于多种任务,可做为Low-level的基线模型;
  2. 作者发现与以往基于transformer的方法不同,Swinir不需要比cnn更多的训练数据,收敛速度也更快;
  3. 结构模块化,可以方便调整出不同复杂度的模型;

  1. Liu Z, Lin Y, Cao Y, et al. Swin transformer: Hierarchical vision transformer using shifted windows[C]//Proceedings of the IEEE/CVF international conference on computer vision. 2021: 10012-10022. ↩︎

这篇关于详解SwinIR的论文和代码(SwinIR: Image Restoration Using Swin Transformer)的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/395849

相关文章

详解nginx 中location和 proxy_pass的匹配规则

《详解nginx中location和proxy_pass的匹配规则》location是Nginx中用来匹配客户端请求URI的指令,决定如何处理特定路径的请求,它定义了请求的路由规则,后续的配置(如... 目录location 的作用语法示例:location /www.chinasem.cntestproxy

CSS will-change 属性示例详解

《CSSwill-change属性示例详解》will-change是一个CSS属性,用于告诉浏览器某个元素在未来可能会发生哪些变化,本文给大家介绍CSSwill-change属性详解,感... will-change 是一个 css 属性,用于告诉浏览器某个元素在未来可能会发生哪些变化。这可以帮助浏览器优化

Python基础文件操作方法超详细讲解(详解版)

《Python基础文件操作方法超详细讲解(详解版)》文件就是操作系统为用户或应用程序提供的一个读写硬盘的虚拟单位,文件的核心操作就是读和写,:本文主要介绍Python基础文件操作方法超详细讲解的相... 目录一、文件操作1. 文件打开与关闭1.1 打开文件1.2 关闭文件2. 访问模式及说明二、文件读写1.

详解C++中类的大小决定因数

《详解C++中类的大小决定因数》类的大小受多个因素影响,主要包括成员变量、对齐方式、继承关系、虚函数表等,下面就来介绍一下,具有一定的参考价值,感兴趣的可以了解一下... 目录1. 非静态数据成员示例:2. 数据对齐(Padding)示例:3. 虚函数(vtable 指针)示例:4. 继承普通继承虚继承5.

前端高级CSS用法示例详解

《前端高级CSS用法示例详解》在前端开发中,CSS(层叠样式表)不仅是用来控制网页的外观和布局,更是实现复杂交互和动态效果的关键技术之一,随着前端技术的不断发展,CSS的用法也日益丰富和高级,本文将深... 前端高级css用法在前端开发中,CSS(层叠样式表)不仅是用来控制网页的外观和布局,更是实现复杂交

Linux换行符的使用方法详解

《Linux换行符的使用方法详解》本文介绍了Linux中常用的换行符LF及其在文件中的表示,展示了如何使用sed命令替换换行符,并列举了与换行符处理相关的Linux命令,通过代码讲解的非常详细,需要的... 目录简介检测文件中的换行符使用 cat -A 查看换行符使用 od -c 检查字符换行符格式转换将

springboot循环依赖问题案例代码及解决办法

《springboot循环依赖问题案例代码及解决办法》在SpringBoot中,如果两个或多个Bean之间存在循环依赖(即BeanA依赖BeanB,而BeanB又依赖BeanA),会导致Spring的... 目录1. 什么是循环依赖?2. 循环依赖的场景案例3. 解决循环依赖的常见方法方法 1:使用 @La

使用C#代码在PDF文档中添加、删除和替换图片

《使用C#代码在PDF文档中添加、删除和替换图片》在当今数字化文档处理场景中,动态操作PDF文档中的图像已成为企业级应用开发的核心需求之一,本文将介绍如何在.NET平台使用C#代码在PDF文档中添加、... 目录引言用C#添加图片到PDF文档用C#删除PDF文档中的图片用C#替换PDF文档中的图片引言在当

详解C#如何提取PDF文档中的图片

《详解C#如何提取PDF文档中的图片》提取图片可以将这些图像资源进行单独保存,方便后续在不同的项目中使用,下面我们就来看看如何使用C#通过代码从PDF文档中提取图片吧... 当 PDF 文件中包含有价值的图片,如艺术画作、设计素材、报告图表等,提取图片可以将这些图像资源进行单独保存,方便后续在不同的项目中使

C#使用SQLite进行大数据量高效处理的代码示例

《C#使用SQLite进行大数据量高效处理的代码示例》在软件开发中,高效处理大数据量是一个常见且具有挑战性的任务,SQLite因其零配置、嵌入式、跨平台的特性,成为许多开发者的首选数据库,本文将深入探... 目录前言准备工作数据实体核心技术批量插入:从乌龟到猎豹的蜕变分页查询:加载百万数据异步处理:拒绝界面