Pyramid attention networks for image restoration

paper:https://arxiv.org/abs/2004.13824

code: https://github.com/SHI-Labs/Pyramid-Attention-Networks

1. 基本思想

作者指出,当前基于深度学习的方法只是在单个尺度上利用了self-similarity(do not take full advantage of self-similarities by relying on self-attention neural modules that only process information at the same scale)。在该论文里,作者提出了 pyramid attention module 用于图像修复,可以利用多尺度特征金字塔捕获 long-range 的特征关联。

在传统图像去噪方法中,NLmean、BM3D都是采用自相似性进行图像去噪的佼佼者,图像的自相似性非常重要。同时,作者指出,基于深度学习的方法大多使用non-local方法来获取 long-range 的信息,这样的模块早期被用于识别和分类任务,但是在 low-level 的视觉任务中也被证明非常有效。

但是,直接将 self-attention 应用于图像修复有一定的局限性:

  • multi-scale non-local prior has not been explorered.
  • pixel-wise matching used in the self-attention module is usually noisy for image restoration.

基于上述分析,作者提出了non-local pyramid attention 模块来捕获long-range信息。

Pyramid attention networks for image restoration_第1张图片

如图所示,该方法可以挖掘 multi-scale self-examplars 来提高重建性能,主要有三个优点:

  • 将现有non-local方法的扩展到多尺度空间(现有non-local方法的搜索空间仅限于金字塔的最底端)
  • 不同大小特征之间的 long-range 依赖被较好的建模
  • 可以降低特征中的噪声,进而提升性能。

2. 方法

Pyramid attention networks for image restoration_第2张图片

如图2a所示,non-local attention 在全图进行响应特征平均,可以这样描述:

\[y^i=\frac{1}{\sigma(x)}\sum\phi(x^i, x^j)\theta(x^j) \]

其中\(\phi\) 计算相似性,\(\theta\)表示特征变换。上述公式仅在统一尺度空间进行处理,难以利用跨尺度的特征。

图2b为 scale agnostic attention,是 2a 的一种扩展。给定一个尺度\(s\),如何评价\(x^j\)\(x^j_{\delta(s)}\)之间的相似性,以及如何把\(x^j_{\delta(s)}\)聚合到\(y^i\)是两个关键问题。为了解决这两个问题,作者将\(x^j_{\delta(s)}\)下采样为像素特征\(z^j\),由于需要在整个特征图上搜索,作者对原始输入\(x\)进行下采样得到特征图\(z\),然后相似性可以这样计算:

\[y^i=\frac{1}{\sigma(x,z)}\sum_j\phi(x^i,z^i)\theta(z^i) \]

同时,作者将该计算方式扩展到金字塔形式,金字塔注意力具有跨尺度特征自相似建模能力。

在具体实现时,相似性函数\(\phi\)使用了embeded gaussian。在特征变换函数\(\theta\)方面,使用简单的线性变换。同时,为了改进匹配的鲁棒性,添加了额外的近邻相似约束,类似于经典的非局部均值滤波,添加了一个强约束信息:近邻且相似。这样有助于学习相关特征,同时抑制不相关特征。

Pyramid attention networks for image restoration_第3张图片

图3 为 Pyramid attention 模块的一个示意图,在损失函数方面,作者选择了L1损失。

3. 实验与结论

作者使用了一个5层的金字塔,网络包含80个残差块,提出的PA模块插入在第40个残差块之后,第41个残差块之前。该方法在图像降噪、去马赛克、去压缩伪影、超分辨率方面都进行了实验,效果都比较好。具体细节和ablation study 可以参考论文,不再细说。

该论文使用了一个 simple backbone,然后只加了一个 pyramid attention 模块,就取得了非常好的效果。作者认为,该模块可以在以后的图像修复网络中,作为基础模块来使用。

你可能感兴趣的:(Pyramid attention networks for image restoration)