论文阅读：Omni-Kernel Network for Image Restoration

论文地址：https://ojs.aaai.org/index.php/AAAI/article/view/27907
项目地址：https://github.com/c-yn/OKNet
发表时间：2024

图像恢复的目的是从一个退化的低质量的观测中重建一个高质量的图像。最近，Transformer模型由于其强大的远程依赖建模能力，在图像恢复任务上取得了很好的性能。然而，相对于输入大小的二次增长的复杂性使它们不适用于实际应用。在本文中，我们开发了一个有效的图像恢复卷积网络，通过增强多尺度表示学习。为此，我们提出了一个由三个分支组成，即全局分支、大分支和局部分支，以有效地学习全局到局部的特征表示。具体地说，全局分支通过 dual-domain channel attention和frequency gated机制实现了全局感知场。此外，为了提供多粒度的感受野，大的分支是通过不同形状的深度卷积和异常大的核大小来制定的。此外，我们使用点级深度卷积来补充局部信息。最后，通过将全核模块插入到bottleneck位置，建立了所提出的网络OKNet。大量的实验表明，我们的网络在11个基准数据集上实现了最先进的性能，用于三个具有代表性的图像恢复任务，包括图像去模糊、图像去雨和图像去模糊。

1、Introduction

图像恢复的目的是从低质量的图像中恢复清晰的图像，如雾霾、雪和模糊场景下的图片。为了处理这个长期存在的不适定问题，传统的方法利用各种手工制作的特征和假设来限制解空间。然而，这些方法并不适用于更具挑战性的现实世界场景（Zhang et al. 2022）。

近年来，卷积神经网络（CNNs）通过从收集到的大规模数据集中学习可推广的先验，在图像恢复任务上取得了优于传统算法的性能（Chen et al. 2019）。为了进一步提高性能，我们已经从其他领域开发或借鉴了许多高级功能单元用于图像恢复，如编码器-解码器架构（Cho等人2021）、残差连接（Mao等人2021）和注意机制（Cui等人2023c）。最近，Transformer模型被引入图像恢复，并显著提高了最先进的性能（Zamir et al. 2022）。尽管如此，Transformer中self-attention，相对于输入大小呈二次增长，这使得这些方法不适合实际应用。·

与卷积有限的接受算子不同，Transformer模型进行全局或大的基于窗口的自我注意，使网络能够获得大的感受野。受这一机制的启发，最近一些关于CNN的工作通过设计具有大内核的高效CNN框架来反击，如RepLKNet中的31×31（Ding等人2022）和SLaK中的51×51（Liu等人2023）。在图像恢复的背景下，LKDNet（Luo et al. 2022）将21×21卷积分解为更小的深度卷积和深度扩展卷积，用于图像去模糊。LaKDNet（Ruan et al. 2023）利用大核大小卷积（例如，9×9），然后进行点卷积，以获得大的有效感受野用于图像去模糊。MAN（Wang et al. 2022b）将一个大的核大小卷积分解为三个分量，即深度卷积和深度扩张卷积，以及点向卷积。然而，由这些方法产生的接受域仍然有限，而且它们不能提供全局的感受野。

Transformer的全局感受野给cnn网络设计带来启发，使用大kernel卷积提取大感受野，并对大kernel进行拆解以获取速度上的优势

在本文中，我们利用63×63深度卷积来探索大的核大小卷积对图像恢复的潜力。此外，我们利用基于条带的卷积进一步增强表示学习高质量的图像重建。为了抑制这些大卷积所带来的计算开销，我们只在瓶颈中部署它们。此外，我们利用双域通道注意和频率门控机制来提供全局的感受野。除了追求大感受野外，我们还利用1×1的深度卷积来补充小规模退化的局部信息。最后，通过将上述设计并行组织，形成了所提出的全核模块（OKM），使网络具有处理多尺度退化的能力。

在瓶颈中使用OKM模块，我们的简单卷积网络在11个不同的数据集上为三个具有代表性的图像恢复任务达到了最先进的性能。更具体地说，OKNet比最近的Transformer模型Fourmer（周等人2023）具有3.47 dB PSNR（周等人2023）的提升，如图1 (a).所示对于单图像离焦去模糊，我们的模型在DPDD（Abuolaim和Brown 2020）数据集的组合类别中，比强Transformer模型恢复器（Zamir等人2022）的性能提高了0.2 dB PSNR。此外，所提出的模型还代表了图像筛选任务的强大能力，在CSD（Chen et al. 2021b）数据集上比最近的算法IRNeXt（Cui et al. 2023c）高出0.7 dB PSNR。

综上所述，本文的贡献可以总结如下：

• 本文提出了一种能够有效捕获多尺度感受野用于图像恢复的全核模块，其中大尺度信息通过双域处理和不同形状的大核尺寸深度卷积进行调制。

• 在11个广泛使用的基准数据集上进行的大量实验表明，所提出的模型，即OKNet，在图像去焦去模糊、图像去雾和图像去雨等三个具有代表性的图像恢复任务上取得了最先进的性能。

2、Related Works

2.1 Image Restoration

图像恢复是一个长期存在的问题，其目标是从其退化的对应图像中重建一个干净的图像，在监控、自动驾驶技术、遥感和医疗成像等许多场景中发挥重要作用。由于其高度不适定的特性，许多传统的算法主要是基于假设和手工制作的特征而开发的，这并不适用于更具有挑战性的实际应用。

近年来，深度学习方法通过从大规模数据集学习可推广的先验，取得了明显优于传统竞争对手的性能。这些方法大致可以分为两类：基于cnn的方法和基于Transformer的方法。多年来，基于CNN的方法通过设计或借用其他领域的高级功能单元来主导图像恢复（Cui et al. 2023b；Cui和Knoll 2023）。例如：

• FFA-Net（Qin et al. 2020）利用通道注意和像素注意模块，不均匀地处理不均匀的烟雾分布特征。
• SDWNet（Zou et al. 2021）利用不同扩张速率的多个扩张卷积并行获得大的感受野。
• SFNet（Cui et al. 2023d）使用动态选择频率模块来选择信息最丰富的频率进行恢复。
• MAN（Wang et al. 2022b）通过将一个大的核卷积分解为三种不同的卷积，提出了对大核的关注。
• LKD-Net（Luo等人，2022年）将深度卷积分解为更小的深度卷积和深度扩张卷积。
  我们的方法与上述的图像恢复算法有四个不同： (a)我们探索了异常大的核大小卷积的潜力，即63×63；(b)除了常规的正方形深度卷积外，我们在不同方向使用基于条的版本来提供不同形状的接受毡进行高质量的图像重建；©我们通过双域处理提供全尺寸的接受毡；(d)我们通过非常简单的1×1深度卷积来补充局部信息。仅在瓶颈中部署所提出的方法，我们高效的OKNet可以与最先进的变压器模型相当或更好。

2.2 Large Kernel Network

最近，受到变压器成功背后的合理原因的启发，即远程依赖建模能力，基于cnn的方法通过使用大型核卷积进行反击。例如：

• RepLKNet（Ding et al. 2022）根据一些设计大型卷积的指导方针实现了31×31内核，极大地缩小了cnn和变压器模型之间的性能差距。
• SLaK（Liu et al. 2023）利用稀疏因子分解的51个×51内核来对抗Transformer方法
• LaKDNet（Ruan et al. 2023）通过大核（9×9）深度向卷积和点向卷积的组合来扩大有效的接受场。
• MAN（Wang et al. 2022b）通过将一个大核卷积分解为三种不同的卷积，呈现了大核的关注。
• LKD-Net（Luo等人，2022年）将深度卷积分解为更小的深度卷积和深度扩张卷积。

我们的方法与上述的图像恢复算法有四个不同：

• (a ) 我们探索了异常大的核大小卷积的潜力，即63×63；
• (b ) 除了常规的正方形深度卷积外，我们在不同方向使用基于条的版本来提供不同形状的接受毡进行高质量的图像重建；
• (c ) 我们通过双域处理提供全尺寸的接受毡；
• (d ) 我们通过非常简单的1×1深度卷积来补充局部信息。
  仅在模型bottleneck层中部署所提出的方法，我们高效的OKNet可以与最先进的变压器模型相当或更好。

3、Methodology

在本节中，我们将首先描述OKNet的Pipeline。然后，我们描述了我们的全内核模块（OKM）的体系结构细节。最后，我们给出了在训练阶段所使用的损失函数。

3.1 Overall Pipeline

整个Pipeline如图2 (a).所示如图所示，OKNet采用了一种编码器-解码器架构，它由编码器和解码器阶段的三个尺度组成。ResGroup由多个残差块组成，每个残差块包括两个3×3卷积，与两者之间的GELU（亨德里克斯和Gimpel 2016）非线性。OKM只插入到瓶颈位置，其中特性的分辨率最低，以节省计算开销。
)

给定一个输入退化rgb图像I，我们利用3×3卷积将图像投影到大小为C×H×W的嵌入特征中，其中C表示通道数，H×W指定像素的空间位置。接下来，将得到的特征输入编码器阶段，以提取深度表示。降采样操作是通过一个串行卷积（kernel=3，strip=2）来实现的，它在减少空间维度的同时扩展了通道数。在被所提出的OKM处理后，这些特征通过解码器网络来恢复高分辨率的表示。在这个过程中，解码器特性与编码器特性连接起来以帮助恢复，然后进行1×1卷积以减少一半的通道。上采样层是通过一个转置卷积（kernel=4，strip=2）来完成的，以扩大空间维度并将通道数量减半。最后，使用3×3卷积得到学习到的残差图像，并加入输入图像产生恢复图像。接下来，我们将详细介绍所提出的OKM。

3.2 Omni-Kernel Module (OKM)

OKM的原理图如图2 (b).所示给定输入特征X，经过1×1卷积处理后，将特征输入三个分支，即局部分支、大分支和全局分支，以增强多尺度表示。三个分支的结果通过相加融合，并通过另一个1×1卷积调制。在下面，我们将介绍每个分支的内部组件。

Large Branch 在这个分支中，我们应用一个低计算量的核大小K×K的深度卷积来追求大的接受毡。除了基于条带的自注意启发的常规深度卷积（Dong等人2022；Tsai等人2022；Li等人2023），我们还使用1×K和K×1平行的深度卷积来获取条形上下文信息。为了避免引入由于较大的核大小卷积而导致的大量计算开销，我们将模块放在瓶颈位置。然后，我们探索了通过逐步增加k来使用极大的卷积进行图像恢复的可能性。实验结果如图3所示。一般来说，峰值信噪比（PSNR）和结构相似性指数（SSIM）指标随着核大小从K = 3扩大到K = 63而增加。我们的模块的放置位置允许我们采用一个异常大的内核大小来捕获大规模的不同接受毡的形状，引入的参数少，复杂度低。为了获得更好的性能，我们将在大分支中选择K = 63。

Global Branch OKNet主要在crop的3×256×256图像块上进行训练，瓶颈特征的空间大小为64×64，因此在大分支中采用最大的奇核大小。然而，在推理阶段，输入的退化图像要比训练size要大得多，一个63×63内核不能覆盖全局域。为了缓解这个问题，我们采用双域处理，在全局分支中增加了全局建模能力。具体来说，全局分支由一个双域通道注意模块（DCAM）和一个基于频率的空间注意模块（FSAM）组成。接下来，我们依次给出这两个模块。

其中F和IF分别为快速傅里叶变压器及其逆运算；$X_{FCA}$、$W_{1\times 1}$和GAP分别表示FCA的输出，一个1×1的卷积层和全局平均池化；⊗表示元素级乘法运算。利用傅里叶处理，根据谱卷积定理有效地再现了全局特征。在光谱域进行全局调制后，得到的特征进一步输入空间通道注意模块（SCA），可以正式表示为：

其中，$X_{DCAM}$是DCAM的输出。DCAM只在通道级粗粒度上增强了双域特性。然后，我们在空间维度中应用基于频率的注意模块，在非粒度水平上重构频谱，其形式表示为：

其中，$X_{FSAM}$是FSAM的结果。通过这样做，该模型关注信息频率分量的高质量的图像重建。

Local Branch 受局部信息在图像恢复中扮演重要角色的启发（Zamir et al. 2022；Wang et al. 2022c），除了捕获大规模感受野的大型全局分支外，我们还设计了一个非常简单但有效的局部信号调制分支，使用1×1深度卷积层，所示我们在表5中演示了其有效性。可以看下表中的Small Branch（第4行），在引入后几乎没有flop的变化，但psnr与ssim都有轻微提升。

3.3 Loss Function

为了恢复真实的高质量图像，一个直接的方法是使预测图像的内容更接近GT图像：

其中，Iˆ和Y分别表示预测的图像和地GT。除了空间域对齐外，所提出的网络还促进了频率信号的学习。因此，我们另外将频域L1损失（Cho et al. 2021）应用于训练：

最后，总体损失函数为：

4、Experiments

在本节中，我们在11个不同的基准数据集上进行了实验，以证明我们的网络对三个具有代表性的图像恢复任务的有效性：图像去模糊、图像离焦去模糊和图像去筛选。在表格中，最好的和第二好的结果用粗体和下划线标记。

4.1 Implementation Details

我们为不同的数据集单独训练模型。根据任务的复杂度，我们通过在每个ResGroup中设置不同数量的残差块（图2 (e)）来扩展OKNet，即N = 4用于去模糊和筛选，N = 16用于去模糊。此外，我们提供了一个小版本的图像去模糊处理，称为OKNet-S，通过设置N = 1来更好地与最近的算法进行比较（Zhou et al. 2023）。除非另有说明，否则将采用以下超参数:

• 使用Adam优化器和β1 = 0.9和β2 = 0.999进行训练
• 批处理大小被设置为8
• 学习速率最初设置为2e−4
• 使用余弦退火衰减策略逐渐降低到1e−6
• 数据增强，大小为256×256的裁剪斑块被随机水平固定，概率为0.5。
• 在256×256的尺寸上计算flop。

4.2 Experimental Results

Image Dehazing Results 我们对三种数据集进行除雾实验：日间合成数据集（RESIDE（李等2018))、日间真实数据集（Dense-Haze（2019）、NH-HAZE（安、安提和雾2020)、 O-Haze（安等2018b)和I-Haze（Ancuti等2018a）)和夜间数据集(NHR (Zhang等202020））。对于日间数据集，我们将我们的模型与9个具有代表性的最先进的模型进行了比较。: GridDehazeNet (Liu et al. 2019), MSBDN (Dong et al.
2020), FFA-Net (Qin et al. 2020), PMNet (Ye et al. 2022),MAXIM-2S (Tu et al. 2022), DeHamer (Guo et al. 2022),SDCE (Zhu et al. 2023), DehazeFormer-L (Song et al.2023), and Fourmer (Zhou et al. 2023).

结果如表1所示

此外，我们的OKNet-S在sourmer上获得了0.27 dB PSNR，流量减少了13%。s-室内和s-t-室外的视觉比较如图4所示。

我们进一步提出比较夜间图像脱雾数据集NHR（张等2020）6最先进的方法：GS（李，谭和布朗2015），MRPF（张等2017），MRP（张等2017），OSFD（张等2020），HCD（王等2022），和焦点Net（崔等2023）。表2显示，我们的方法比最近的FocalNet要好2.57 dB PSNR和0.01 SSIM。

图5说明了我们的网络产生的结果更接近GT。

Image Defocus Deblurring Results 我们使用广泛使用的DPDD数据集验证单图像离焦去模糊网络的有效性，并比较7种代表性算法：DPDNet、KPAC、DRBNet、IFAN、MDP等。结果如表3所示。如上所示，我们的模型在大多数指标上都比其他方法获得了更好的性能。具体地说，OKNet在室外场景中比基于强变压器的模型修复器显著增加了0.27 dB PSNR。此外，与同样使用大核卷积的LaKDNet相比，我们的模型在参数减少21%的组合类别上获得了0.03 dB PSNR和0.002 SSIM的性能提高，如图1 (b).所示可视化的比较如图6所示。我们的方法比其他竞争对手产生了更清晰、更符合人类感官的结果。

Image Desnowing Results 我们在三个广泛用于图像筛选的数据集上评估了所提出的模型，包括Snow100K（Liu等人2018）、SRRS（Chen等人2020）和CSD（Chen等人2021b）。然后，我们将结果与8种最先进的算法进行比较。从表4中可以看出，该网络具有很强的Desnowing 能力。

具体地说，OKNet在Snow100K数据集上比最近的算法IRNeXt高出0.14 dB PSNR。在更复杂的CSD数据集上，优势变得更大，我们的方法表明其有效性。图7显示，我们的结果在视觉上比其他结果更令人满意。

4.3 Ablation Studies

我们通过对OKNet-S的RESIDE-Indoor数据集进行300个epoch的消融研究，并对SOTS-Indoo进行评估（Li等人2018年）。基线模型是通过从我们的模型中去除OKM而获得的。

我们逐步将设计的大分支、小分支和全局分支添加到基线模型中。结果如表5所示。基线模型达到31.32 dB PSNR。大卷积比基线显著增加了3.75 dB，而基于条带的卷积进一步将性能提高到35.29 dB，证明了捕获不同形状的感受野的有效性。这个极其简单的小分支通过增强局部信息，将精度提高到35.48 dB。最后，我们研究了全局分支中单个组件的效率。与基线模型相比，FCA、SCA和FSAM实现了1.52 dB、1.03 dB和2 dB的性能提高。FCA和SCA的组合（#8）比仅使用FCA（#6）产生的得分更高，这表明我们的设计具有兼容性。配备FSAM后，完整模型的效果最好，比基线模型高5.16 dB。

5、Conclusion

在本文中，我们提出了一个有效的卷积网络，称为OKNet，它能够捕获多尺度的感受野。核心组件OKM由三个分支组成，分别用于建模局部、大型和全局依赖关系。大分支是通过探索异常大的convKxK和基于条带的convKx1或conv1xK来设计的。该新的全局分支利用双域通道注意和基于频率的空间注意来调制全局表示。此外，非常轻量级的局部分支给模型带来了局部性。将简单而有效的OKM插入到瓶颈中，OKNet在11个不同的数据集上为三个图像恢复任务实现了最先进的性能。