【论文阅读】图像修复 - NAFNet

Simple Baselines for Image Restoration

https://arxiv.org/abs/2204.04676

亮点

---

1. 网络架构简单，计算量相对较低
2. 实验证明了网络中的激活函数可用乘法代替或移除
3. 效果达到了 SOTA，计算量也低

GoPro图像去模糊任务上，较之前 SOTA 的 PSNR 高 0.38db，但计算量降低了 91.6%

SIDD图像去噪任务上，较之前 SOTA 的 PSNR 高 0.28db，但计算量降低了 50%

如何做

---

论文中构建简单网络架构可以用下面三点来概括

• 使用带 short-cut 的 UNet 作为网络的 Backbone
• 去除 Gated Linear Units 中的 tanh
• 去除 Channel Attention 中的 Sigmoid，ReLU

UNet

为了减少模型每个模块间复杂的交互，网络架构为带 short-cut 的 UNet，同时还引入了 LN 操作，因为在 Transformer 中验证了它能使训练更平滑

Gated Linear Units

多篇相关工作中同时用到了 GLU (Gated Linear Units)并且声尘是有收益，因此 NAFNet 也引入了这个模块，其中 X 为 feature map，f 和 g 为linear transformers，σ 为激活函数

但加入 GLU 会增加计算量，大头主要是在激活函数这块，因此论文提出了一个观点，把激活函数去掉的 GELU 还是具备非线性拟合能力的

基于此提出了 Simple Gate

class SimpleGate(nn.Module):
    def forward(self, x):
        x1, x2 = x.chunk(2, dim=1)
        return x1 * x2
1
2
3
4

从计算上来说，x1 和 x2 在训练时随参数更新而变化，说具备非线性拟合能力也说的通，论文实验发现在图像去噪和去模糊任务上分别有 0.08db 和 0.41db 的收益

Channel Attention

同样，加入了通道注意力 CA 组件，对效果也是有收益的，但论文这里对传统注意力组件做了简化，没错，也是去除激活函数，同样的也命名为 Simple Channel Attention，经实验，PSNR 值在 SIDD 上有 0.14 db 增益，而在 GoPro 上有 0.24db 增益

# Simplified Channel Attention
self.sca = nn.Sequential(
     nn.AdaptiveAvgPool2d(1),
     nn.Conv2d(in_channels=dw_channel // 2, out_channels=dw_channel // 2, kernel_size=1, padding=0, stride=1,
     groups=1, bias=True),
)

1
2
3
4
5
6
7
8

Metric

---

从计算量，PSNR 两个角度衡量该算法在不同图像处理任务下的效果

更多对比细节，请看原文