【论文阅读】Uformer:AGeneralU-ShapedTransformerforImageRestoration

博客 动态
0 315
羽尘
羽尘 2023-06-09 00:10:37
悬赏:0 积分 收藏

【论文阅读】Uformer:A General U-Shaped Transformer for Image Restoration

??前言

  • ??博客主页:??睡晚不猿序程??
  • ?首发时间:2023.6.8
  • ?最近更新时间:2023.6.8
  • ??本文由 睡晚不猿序程 原创
  • ??作者是蒻蒟本蒟,如果文章里有任何错误或者表述不清,请 tt 我,万分感谢!orz


1. 内容简介

论文标题:Uformer: A General U-Shaped Transformer for Image Restoration

发布于:CVPR 2021

自己认为的关键词:Transformer、Unet

是否开源?https://github.com/ZhendongWang6/Uformer


2. 论文浏览

论文动机

  1. 目前对于图像修复(image restoration) 的 SOTA 模型都是基于 ConvNet 的,其把握全局信息的能力不足
  2. 一些针对上面的问题的改进方式引入注意力层,但是由于其平方级别的复杂度导致只能使用较少层数

本文工作

  1. Uformer:搭建了分层级的 encoder-decoder 模型(类似 U-Net)

  2. LeWin:基于 Swin Transformer 的滑动窗口自注意力上在 FFN 添加 DWConv

  3. 可学习的多尺度修复模块:用于调整 decoder 部分的特征,使其自适应不同噪声

    一个和 Token 维度相同的可学习参数

完成效果:在去噪、去雨、去模糊等 low level 任务效果显著

3. 图片、表格浏览

图一

image-20230608233858657

模型性能图,看起来在相同计算量下比 Unet 性能要好

奇怪的是怎么没有和 Swin Transformer 碰碰

图二

image-20230608233918473

模型架构图,可以看到组成了一个 Unet 架构,在 decoder 分支使用了 Modulators 调整特征的表示

图三

image-20230608233929682

局部增强的 FFN,在 1x1 卷积中间引入了 DWConv

感觉可以像 ConvNeXt 一样挪到 1x1 卷积前面,可以进一步降低参数量

图四

image-20230608234025369

Modulator 的效果,看起来确实不错

Token 的表示确实很轻松就可以引入某些信息

图五

image-20230608234035047

模型去噪效果


4. 引言浏览

本文重心:在不同分辨率上使用自注意力,从而恢复图像信息——提出了 Uformer,基于 U-Net,但是保持了他的 味道 ,只是改变了 backbone(跳跃连接啥的没有变)

两个设计

  1. Locally enhanced Window:字如其名
    特征图分辨率改变,但是窗口大小不变
    在 Transformer FFN 的两个全连接层之间添加 depth-wise 卷积层
  2. learnable multi-scale restoration modulator:用于学习图像的 退化方式
    作为一种多尺度的空间偏置,来在 decoder 中调整图像特征——可以可学习的,基于窗口的 tensor 将会直接加到特征中去,来调整这些特征,来为重建提供更多信息

自由阅读


5. 方法

5.1 模型结构

image-20230608234047932

encoder

  1. 使用 3x3 LeakyReLU 卷积对输入图像进行处理,得到(C,H,W)大小图像,通过此抽取底层特征

  2. 接着就是和 UNet 一样,要经过 K 个 encoder stages,每个 stage 由一个 LeWin block 和一个下采样层组成

    具体而言,LeWin block 输入输出分辨率相同,下采样层先把图片转化为 2D 的(从 Token 中恢复图像)然后使用 4x4,stride=2 的卷积降低分辨率,翻倍特征通道数

  3. encoder 的最后添加上一个 bottleneck stage,由一系列 LeWin Transformer block 堆叠而成
    在这里可以使用全局自注意力,相比之前的窗口自注意力可以有更强的建模能力

decoder

  1. 同样 K 个 stage,包含若干个 Transformer blcok 以及一个上采样层
    具体而言,使用 2x2 stride=2 的转置卷积来做上采样,翻倍分辨率,减半特征通道数量

  2. 接着,和之前的部分做级联,然后一直计算下去

  3. 经过了 K 个 stage,将其重新转化为 2D 表示,并使用一个 3x3 卷积将其特征通道数重新变为 3,然后进行一个残差连接得到输出

    这里的上下采样选择转为图像表示再使用卷积进行操作

LeWin Transformer Block

ViT 两个缺陷:

  1. 全局自注意力
  2. 难以把握 local dependencies

进行了改进,结构如下

image-20230608234109931

两个核心设计:

  1. 不重叠的基于窗口的自注意力(W-MSA)
  2. 局部强化前向网络(LeFF)

LeFF

image-20230608234118930

  1. 首先进行线性映射
  2. 转化为图像表示,进行 3x3 卷积
  3. 转为 token,然后映射为原本的维度,进行残差连接

Multi-Scale Restoration Modulator

这里说到,因为不同的图像退化方式会有不同的扰动模式,为了让 Uformer 可以解决这些不同的扰动,使用了一个轻量化的恢复模块,它可以标定输出的特征,鼓励模型使用更多的细节来恢复图像

image-20230608234132716

在每一个块中,这个调制器是一个 tensor ,大小为(M,M,C),用法和偏置一样,在计算自注意力之前直接加到每一个窗口中就可以了。

在去模糊和图像去噪方面可以改善效果,尤其是在去运动模糊方面和噪声方面可以得到更好的结果

这里说,和 StyleGAN 中添加随机噪声的思路是一致的,所以这个 StyleGAN 之后也可以去看一下

6. 实验

6.1 实验设置

损失函数

\[l(I',\hat I)=\sqrt{||I'-\hat I||^2 +\epsilon^2} \]

这里是选择了一个 Charbonnier Loss 作为损失函数

优化器:Adam(0.9, 0.999), weight decay 0.02; lr 2e-4;使用余弦衰减

数据增强:水平翻转,旋转 90,180,270

模型架构

image-20230608234146560

窗口大小 4x4

评价指标

  • PSNR
  • SSIM

使用 YUV 的 Y 通道进行计算

6.2 实验结果

Real Noise Removal

第一项任务,去噪

数据集:SIDD,DND

对比了 8 个模型,拿到了最佳成绩

Motion blur removal

任务二:去除动态模糊

数据集:GoPro(训练),HIDE(合成的数据集),GoPro(测试集),RealBlur-R/J

在 GoPro 上对之前的 SOTA 方法实现了全面的超越(PSNR 多了 0.3 个点,SSIM 多了 0.01)其他的数据集也都比之前的 SOTA 要好

defocus Blur Removal

任务三:去失焦

数据集:DPD

全面超越,PSNR 多了好几个点!直接按照一位数算的

Real Rain Removal

离谱啦,多了 3 个点

6.3 消融实验

1. Transformer vs convolution

将其中的 Tr 替换为 CNN(ResBlock),设计了三个不同大小的 UNet

image-20230608234158552

效果不错,参数量更小的情况下能得到更好的效果,但是我觉得差距好像也不是很大?

哦但是那个 B 模型就差距比较大了,多了 0.1 个点

2. Hierarchical structure vs single scale

使用 ViT-based 架构,单一尺度,用于图像去噪,在开始和结束的时候使用两层卷积来提取特征以及还原图像,其中使用 12 层 Tr block,隐藏特征维度为 256,patch 大小 16x16,得到的效果不好

ViT 的单一尺度对这种任务效果肯定不好,ViT 似乎是一个 近视眼 ,16x16 的 patch 划分让他具有了这种特性,可以看成是低通滤波器

3. Where to enhance locality

对比了是否进行局部性强化的效果

image-20230608234204412

可以看到,在自注意力中引入局部强化会导致效果变差,但是如果在 FFN 中引入局部强化会得到一个更好的效果

对 Token 做卷积可以看成是一个超大感受野的扩张卷积

4. Effect of the multi-scale restoration modulator

image-20230608234214648

效果显著,如果看他的效果对比,可以看的更为明显

image-20230608234221937

Is Window Shift Iportant?

附录内容

使用滑动窗口,带来了微小的提升

但是根据我自己跑代码的情况,我使用了滑动窗口带来的提升挺大的(2-3 个点 PSNR),约等于进行了一次感受野的扩大

6. 总结、预告

6.1 总结

使用 Transformer 作为 backbone 的 U-Net 网络,具有两个创新点

  1. 基于窗口的自注意力,自注意力卷积联合使用,随着网络的加深模型可以获得更大的 感受野
  2. 可学习的 multi-scale restoration modulator,仅仅使用一个加法取得了一个不错的效果
  3. 卷积和 MLP 相结合,构成了一个 局部性强化 方法,更好的应用图像的局部信息

6.2 预告

打算之后更新一下关于 Transformer 的论文阅读了

好久没有整理自己的笔记然后更新了iai,屯了一大堆的论文笔记

posted @ 2023-06-08 23:46  睡晚不猿序程  阅读(0)  评论(0编辑  收藏  举报
回帖
    羽尘

    羽尘 (王者 段位)

    2335 积分 (2)粉丝 (11)源码

     

    温馨提示

    亦奇源码

    最新会员