《MPRnet》学习笔记

paper：2102.02808

GitHub：swz30/MPRNet: [CVPR 2021] Multi-Stage Progressive Image Restoration. SOTA results for Image deblurring, deraining, and denoising.

摘要

1、介绍

2、相关工作

2.1 单阶段方法

2.2 多阶段方法

2.3 注意力机制

3、方法

3.1 整体框架

3.2 互补特征处理

3.2.1 编码器-解码器子网络

3.2.2 原始分辨率子网络

3.3 跨阶段特征融合（CSFF）

3.4 有监督注意力模块（SAM）

4、实验

4.1. 数据集与评估协议

4.2. 实现细节

4.3. 图像去雨结果

编辑 4.4. 图像去模糊结果

4.5. 图像去噪结果

4.6. 消融研究

摘要

图像恢复任务需要在恢复图像时实现空间细节与高级上下文信息之间的复杂平衡。本文提出了一种新颖的协同设计，能够在恢复过程中优化平衡这些相互竞争的目标。我们的主要提议是一个多阶段架构，通过逐步学习退化输入的恢复函数，将整体恢复过程分解为更易管理的步骤。具体来说，我们的模型首先使用编码器-解码器架构学习上下文化特征，随后将这些特征与一个保持局部信息的高分辨率分支相结合。在每个阶段，我们引入了一种新颖的每像素自适应设计，利用现场监督注意力来重新加权局部特征。多阶段架构的一个关键成分是不同阶段之间的信息交换。为此，我们提出了一种双重方法，其中信息不仅从早期阶段向后期阶段顺序交换，而且特征处理模块之间还存在横向连接，以避免信息丢失。最终，紧密相连的多阶段架构——称为MPRNet，在包括图像去雨、去模糊和去噪等任务的十个数据集上，展现出了强劲的性能提升。

1、介绍

图像恢复是从退化版本中恢复干净图像的任务。退化的典型例子包括噪声、模糊、雨水、雾霾等。由于存在无限多个可行的解决方案，这个问题是一个高度不适定的问题。为了将解决方案空间限制为有效/自然的图像，现有的恢复技术明确地使用了通过经验观察手工设计的图像先验。然而，设计这些先验是一个具有挑战性的任务，并且通常不具有可扩展性。为了缓解这个问题，最近的最先进方法采用了卷积神经网络（CNNs），通过从大规模数据中捕获自然图像的统计特性，隐式学习更通用的先验。

基于CNN的方法相较于其他方法的性能提升，主要归因于其模型设计。已经开发了许多图像恢复的网络模块和功能单元，包括递归残差学习、膨胀卷积、注意力机制、密集连接、编码器-解码器和生成模型。然而，几乎所有这些用于低层视觉问题的模型都是基于单阶段设计。相比之下，多阶段网络在高层视觉问题（如姿态估计、场景解析和动作分割）中被证明比单阶段方法更有效。最近，也有一些尝试将多阶段设计应用于图像去模糊和去雨。我们分析了这些方法，识别出限制它们性能的架构瓶颈。首先，现有的多阶段技术要么使用编码器-解码器架构，该架构有效地编码了广泛的上下文信息，但在保持空间图像细节方面不可靠，要么使用单尺度管道，提供空间上准确但语义上不太可靠的输出。然而，我们证明，结合这两种设计选择的多阶段架构是图像恢复的有效方案。其次，我们证明，简单地将一个阶段的输出传递给下一个阶段会导致次优结果。第三，在每个阶段提供地面真实监督以进行渐进恢复是非常重要的。最后，在多阶段处理中，需要有一个机制来传播从早期阶段到后期阶段的中间特征，以保留来自编码器-解码器分支的上下文化特征。

我们提出了一种多阶段渐进式图像恢复架构，称为MPRNet，具有几个关键组件：

1). 多框架设计：早期阶段使用编码器-解码器来学习多尺度上下文信息，而最后一个阶段在原始图像分辨率下操作，以保持细致的空间细节。

2). 每两个阶段之间插入一个监督注意力模块（SAM），以实现渐进学习。在地面真实图像的指导下，该模块利用前一阶段的预测计算注意力图，这些图再用于在传递到下一个阶段之前精细调整前一阶段的特征。

3). 添加了跨阶段特征融合（CSFF）机制，帮助将多尺度上下文化特征从早期阶段传播到后期阶段。此外，这种方法简化了阶段之间的信息流动，有助于稳定多阶段网络的优化。

该工作的主要贡献包括：

提出了一种新颖的多阶段方法，能够生成具有上下文化和空间精确的输出。由于其多阶段的性质，我们的框架将复杂的图像恢复任务分解为子任务，逐步恢复退化图像。

一种有效的监督注意力模块，充分利用每个阶段恢复的图像，在精细调整传入特征时提供帮助，然后再进一步传播这些特征。

一种跨阶段聚合多尺度特征的策略。

我们通过在十个合成和真实世界的数据集上进行实验，展示了MPRNet的有效性，并在图像去雨、去模糊和去噪等恢复任务中设置了新的最先进性能，同时保持较低的复杂度。此外，我们提供了详细的消融实验、定性结果和泛化测试。

2、相关工作

近年来，图像捕捉领域经历了一次范式转变，从高端单反相机（DSLR）转向智能手机相机。然而，使用智能手机相机拍摄高质量图像仍然具有挑战性。图像退化常常出现在图像中，这可能是由于相机的限制和/或不利的环境条件。早期的恢复方法基于总变差（Total Variation）、稀疏编码（Sparse Coding）、自相似性（Self-Similarity）、梯度先验（Gradient Prior）等。最近，基于CNN的恢复方法已经达到了最先进的结果。从架构设计的角度来看，这些方法大体上可以分为单阶段和多阶段两种。

2.1 单阶段方法

目前，大多数图像恢复方法都基于单阶段设计，且架构组件通常基于那些为高层视觉任务开发的模块。例如，残差学习（Residual Learning）已被用于图像去噪、图像去模糊和图像去雨。类似地，为了提取多尺度信息，通常使用编码器-解码器膨胀卷积模型。其他单阶段方法结合了密集连接（Dense Connections）。

2.2 多阶段方法

这些方法旨在通过在每个阶段采用轻量级子网络，以渐进的方式恢复干净图像。这种设计非常有效，因为它将复杂的图像恢复任务分解为较小、更易处理的子任务。然而，常见的做法是为每个阶段使用相同的子网络，这可能会导致次优结果。

2.3 注意力机制

受到其在高层任务（如图像分类、分割和检测）中的成功启发，注意力模块也被用于低层视觉任务。例如，图像去雨、去模糊、超分辨率和去噪的方法中都广泛应用了注意力机制。其主要思想是捕捉沿空间维度、通道维度或两者的长程依赖关系。

3、方法

3.1 整体框架

我们提出的图像恢复框架，如图所示，包含三个阶段，用于逐步恢复图像。

前两个阶段基于编码器-解码器子网络，学习广泛的上下文信息，因为它们具有较大的感受野。由于图像恢复是一个位置敏感的任务（需要输入和输出之间的逐像素对应关系），因此最后一个阶段采用一个操作在原始输入图像分辨率上工作的子网络（没有进行下采样操作），从而在最终输出图像中保留所需的精细纹理。

我们不仅仅是简单地级联多个阶段，而是在线性每两个阶段之间加入了一个监督注意力模块。在地面真实图像的监督下，我们的模块在将前一阶段的特征图传递给下一阶段之前对其进行重新缩放。此外，我们引入了跨阶段特征融合机制，其中早期子网络的中间多尺度上下文特征有助于巩固后期子网络的中间特征。

尽管 MPRNet 堆叠了多个阶段，但每个阶段都可以访问输入图像。与最近的恢复方法类似，我们对输入图像适配了多patch层次结构，将图像分成不重叠的patch：第一个阶段为四个patch，第二个阶段为两个patch，最后一个阶段为原始图像，如图所示。

在任意给定阶段 $S$ 时，我们的模型并非直接预测恢复图像 $X_S$ ，而是预测一个残差图像 $R_S$ ，降解的输入图像 $I$ 与其相加得到：

$X_S = I + R_S$

我们使用以下损失函数对MPRNet进行端到端优化：

$L = \sum_{S=1}^{3} \left[ L_{\text{char}}(X_S, Y) + L_{\text{edge}}(X_S, Y) \right]$

其中， $Y$ 表示真实图像， $L_{\text{char}}$ 为Charbonnier损失：

$L_{\text{char}} = \sum_{i} \sqrt{(X_S[i] - Y[i])^2 + \epsilon^2}$

其中常数 $\epsilon =10^{-3}$ 在所有实验中设定为固定值。

此外， $L_{\text{edge}}$ 是边缘损失，定义为：

$L_{\text{edge}} = \sqrt{\left \| \Delta (X_S)-\Delta (Y) \right \|^2+\epsilon ^2}$

其中， $\Delta$ 表示拉普拉斯算子。参数 λ 控制了两个损失项的相对重要性，实验中我们将其设置为0.05。

3.2 互补特征处理

现有的单阶段CNN图像恢复方法通常采用以下两种架构设计之一：1）编码器-解码器架构，或者2）单尺度特征管道。编码器-解码器网络首先将输入逐渐映射到低分辨率表示，然后通过逐步反向映射恢复原始分辨率。虽然这些模型能有效编码多尺度信息，但由于反复使用下采样操作，它们容易牺牲空间细节。与此相反，操作在单尺度特征管道上的方法可靠地生成具有细腻空间细节的图像。然而，由于感受野有限，它们的输出在语义上较为薄弱。这表明，前述架构设计选择的固有局限性：它们要么能够生成空间上准确的输出，要么能够生成语义上可靠的输出，但不能同时具备两者。

为了充分利用两种设计的优点，我们提出了一种多阶段框架，其中早期阶段结合了编码器-解码器网络，最后阶段采用操作在原始输入分辨率上的网络。

3.2.1 编码器-解码器子网络

图展示了我们的编码器-解码器子网络，它基于标准的U-Net架构，具有以下几个组件。首先，我们添加了通道注意力模块（CABs），用于在每个尺度上提取特征。其次，U-Net跳跃连接中的特征图也会通过CAB进行处理。最后，我们没有使用反卷积（转置卷积）来增加解码器中特征的空间分辨率，而是使用双线性上采样（bilinear upsampling）和卷积层。这有助于减少由于转置卷积常常产生的棋盘状伪影。

3.2.2 原始分辨率子网络

为了从输入图像到输出图像保留精细的细节，我们在最后阶段引入了原始分辨率子网络（ORSNet）。ORSNet不进行任何下采样操作，并生成空间上丰富的高分辨率特征。它由多个原始分辨率块（ORBs）组成，每个块进一步包含CAB。

3.3 跨阶段特征融合（CSFF）

在我们的框架中，我们在两个编码器-解码器之间以及编码器-解码器与原始分辨率子网络（ORSNet）之间引入了跨阶段特征融合（CSFF）模块。需要注意的是，一个阶段的特征在传播到下一个阶段进行聚合之前，会先通过1x1卷积进行精炼。提出的CSFF具有几个优点。首先，它通过减少编码器-解码器中反复使用上采样和下采样操作导致的信息丢失，使得网络不容易受到影响。其次，一个阶段的多尺度特征有助于丰富下一个阶段的特征。最后，网络优化过程变得更加稳定，因为它缓解了信息流的阻塞，从而允许在整体架构中添加多个阶段。

3.4 有监督注意力模块（SAM）

近期的一些多阶段图像恢复网络直接在每个阶段预测图像，并将其传递到下一个连续阶段。与此不同，我们在每两个阶段之间引入了有监督注意力模块（SAM），它有助于显著提高性能。SAM的示意图如图所示。

其贡献有两个方面。首先，它提供了有助于每个阶段图像恢复的真实值信号。其次，在局部监督预测的帮助下，我们生成注意力图，以抑制当前阶段中不太有用的特征，只允许有用的特征传播到下一个阶段。

如图4所示，SAM首先接收来自前一个阶段的输入特征 $F_{in}\in \mathbb{R}^{{H \times W \times C}}$ ，并通过简单的1x1卷积生成一个残差图像 $R_{S}\in \mathbb{R}^{{H \times W \times 3}}$ 。然后将残差图像加到降质输入图像 $I$ 上，以得到恢复图像 $X_{S}\in \mathbb{R}^{{H \times W \times 3}}$ 。对于这个预测图像 $X_{S}$ ，我们提供真值。接下来，从图像 $X_{S}$ 中通过1x1卷积和Sigmoid激活生成逐像素的注意力掩膜 $M\in \mathbb{R}^{H \times W \times C}$ 。这些掩膜随后用于重新校准经过1x1卷积转换后的局部特征 $F_{in}$ ，得到注意力引导的特征，这些特征会加到恒等映射路径上。最后，由SAM产生的增强特征表示 $F_{out}$ 被传递到下一个阶段进行进一步处理。

4、实验

我们对多种图像恢复任务评估了我们的方法，包括：(a) 图像去雨、(b) 图像去模糊、(c) 图像去噪，使用了10个不同的数据集。

4.1. 数据集与评估协议

定量比较使用PSNR和SSIM指标进行。我们报告每个方法相对于最优方法的误差减少，通过将PSNR转化为RMSE和SSIM转化为DSSIM。用于训练和测试的数据集在表1中总结，接下来将描述。

图像去雨：使用与最近最优的图像去雨方法相同的实验设置，我们在13712对清晰-降雨图像对上进行训练，这些图像来自多个数据集，如表1所示。使用这个单一训练的模型，我们在多个测试集上进行评估，包括Rain100H 、Rain100L 、Test100、Test2800 和Test1200。

图像去模糊：我们使用GoPro数据集，该数据集包含2103对图像用于训练，1111对图像用于评估。此外，为了展示模型的泛化能力，我们将GoPro训练的模型直接应用于HIDE和RealBlur数据集的测试图像。HIDE数据集专门为人类感知运动去模糊收集，测试集包含2025张图像。虽然GoPro和HIDE数据集是合成生成的，但RealBlur数据集中的图像对是在真实条件下拍摄的。RealBlur数据集有两个子集：（1）RealBlur-J是由相机JPEG输出生成的；（2）RealBlur-R是通过对RAW图像应用白平衡、去马赛克和去噪操作离线生成的。

图像去噪：为了训练我们的图像去噪模型，我们使用了SIDD数据集中的320张高分辨率图像。评估在SIDD数据集中的1280个验证补丁和DND基准数据集中的1000个补丁上进行。这些测试补丁是从原始图像中提取的。SIDD和DND数据集均由真实图像构成。

4.2. 实现细节

我们的MPRNet是端到端可训练的，并且无需预训练。我们为三种不同任务训练了单独的模型。我们在编码器-解码器的每个尺度上使用了2个CAB，并且在下采样时使用了2x2的最大池化，步幅为2。在最后阶段，我们使用了包含3个ORB的ORSNet，每个ORB进一步使用8个CAB。根据任务的复杂性，我们通过设置不同的通道数量来调整网络宽度，训练时设置为40，去噪时为80，去模糊时为96。网络在256x256的patch上训练，批量大小为16，训练迭代次数为 $4\times10^5$ 。为了进行数据增强，随机应用水平和垂直翻转。我们使用Adam优化器，初始学习率为 $2 \times10^-4$ ，采用余弦退火策略将学习率逐步降至 $1\times10^-6$ 。

4.3. 图像去雨结果

在图像去雨任务中，与之前的工作一致，我们计算使用Y通道（YCbCr色彩空间）进行的图像质量评分。表2显示，我们的方法在所有五个数据集上都明显超越了现有的最优方法， consistently achieving better PSNR/SSIM scores。与最近最优算法MSPFN相比，我们的模型在所有数据集上的平均性能提升了1.98dB，表明误差减少了20%。在某些数据集上，提升达到4dB，例如Rain100L。此外，我们的模型比MSPFN少了37%的参数，且速度快了24%。图5显示了在具有挑战性的图像上的视觉对比。我们的MPRNet有效去除了不同方向和幅度的雨滴，生成的图像在视觉上令人愉悦，并忠实于地面真值。相比之下，其他方法在结构内容上存在妥协（第一行），引入了伪影（第二行），并未完全去除雨滴（第三行）。

4.4. 图像去模糊结果

我们报告了评估的图像去模糊方法在合成的GoPro和HIDE数据集上的表现，如表3所示。总体而言，我们的模型在其他算法中表现优异。与之前表现最好的技术相比，我们的方法在GoPro数据集上获得了9%的PSNR提升和21%的SSIM提升，在HIDE数据集上则获得了11%和13%的误差减少。值得注意的是，我们的网络仅在GoPro数据集上训练，但在HIDE数据集上达到了最先进的结果（+0.98dB），这证明了其强大的泛化能力。

我们还在真实世界的RealBlur数据集上评估了我们的MPRNet，进行了两种实验设置：1）直接将GoPro训练的模型应用于RealBlur（测试模型对真实图像的泛化能力）；2）在RealBlur数据集上进行训练和测试。表4展示了实验结果。在设置1中，我们的MPRNet相较于DMPHN算法在RealBlur-R子集上获得了0.29dB的性能提升，在RealBlur-J子集上获得了0.28dB的提升。在设置2中，我们在RealBlur-R和RealBlur-J子集上相较于SRN[71]分别获得了0.66dB和0.38dB的提升。

图显示了由评估方法处理的去模糊图像。总体而言，我们模型恢复的图像比其他方法更锐利，更接近地面真值。

4.5. 图像去噪结果

在表5中，我们报告了SIDD和DND数据集上多种图像去噪方法的PSNR/SSIM得分。我们的方法在最先进的方法上取得了显著的提升，即在SIDD上比CycleISP高了0.19dB，在DND上比SADNet高了0.21dB。值得注意的是，DND数据集不包含任何训练图像，公开发布的数据集仅为测试集。在使用SIDD数据集训练的模型在DND基准数据集上的实验结果表明，我们的模型能够很好地泛化到不同的图像领域。

图展示了视觉结果。我们的方法能够去除真实噪声，同时保留图像的结构和纹理细节。相比之下，其他方法恢复的图像要么过于平滑，要么包含伪影和斑驳的纹理。

4.6. 消融研究

在这里，我们展示了消融实验，以分析我们模型中每个组件的贡献。评估在GoPro数据集上进行，去模糊模型使用128x128的图像patch 进行训练，共进行了105次迭代，结果如表6所示。

阶段数：随着阶段数的增加，我们的模型表现更好，这验证了我们多阶段设计的有效性。

子网络选择：由于我们模型的每个阶段可以采用不同的子网络设计，我们测试了不同的选项。结果表明，在早期阶段使用编码器-解码器，在最后阶段使用ORSNet，能够提高性能（297dB），相比之下，使用相同设计的所有阶段（U-Net+U-Net为294dB，ORSNet+ORSNet为295.3dB）表现较差。

SAM和CSFF：通过从最终模型中去除SAM和CSFF，我们展示了这两种机制的有效性。表6显示，当移除SAM时，PSNR从30.49dB降至30.07dB；当移除CSFF时，PSNR从30.49dB降至30.31dB。去除这两个组件会显著降低性能，从30.49dB降至29.86dB。