AMSA-UNet | 基于自注意力的多尺度 U-Net 提升图像去模糊性能 ！

传统的单尺度U-Net在去模糊过程中常常会导致空间信息的丢失，这影响了去模糊的准确性。此外，由于卷积方法在捕捉长距离依赖方面的局限性，恢复图像的质量下降。 为了解决上述问题，提出了一种基于自注意力的非对称多尺度U-Net（AMSA-UNet）来提高去模糊方法在准确性和计算复杂度方面的性能。通过引入多尺度U形结构，网络可以在全局层面关注模糊区域，并在局部层面更好地恢复图像细节。 为了克服传统卷积方法在捕捉信息长距离依赖方面的限制，在主干网络的解码器部分引入了自注意力机制，这显著增加了模型的感受野，使模型能够更多地关注图像的语义信息，从而产生更准确、视觉上更令人满意的去模糊图像。 更重要的是，引入了一种基于频域的计算方法来减少计算量。实验结果表明，与八种优秀的方法相比，所提出的方法在准确性和速度方面都有显著提升。

1 Introduction

早期的去模糊方法主要关注非盲去模糊，恢复已知模糊核的图像。Pan等人[1]通过计算模糊图像中暗通道的稀疏性来准确计算模糊核，以恢复清晰图像。然而，这些传统方法在处理空间变化的模糊方面存在困难，而且通常耗时。

随着深度学习的发展，基于CNN的非盲去模糊方法越来越受到重视。包括基于CNN的模糊核估计和利用估计核的反卷积网络。Sun等人[2]使用CNN增强模糊核中的运动平滑性，而Chakrabarti等人[3]对模糊核系数进行建模以实现精确的反卷积。结合传统方法和深度学习方法改进了去模糊效果，但在处理遮挡、深度变化和模糊核估计中的噪声敏感性方面仍存在挑战，这限制了它们在复杂场景中的有效性。

为了克服非盲去模糊的局限性，已经开发出了一种端到端的卷积神经网络方法。这些方法直接将模糊图像映射到清晰图像，避免了依赖于模糊核估计。Nah等人[4]采用了一种多尺度从粗到精的方法，在动态场景中有效去模糊。Mei等人[5]提出的DeepDeblur模型以端到端的方式去除文档图像中的文本模糊而著称。然而，由于复杂的网络结构，这些方法需要大量的训练时间。DeblurGAN由Kupyn等人[6]提出，它结合了生成对抗网络（GAN）进行去模糊，通过生成器和判别器之间的博弈取得了良好的效果。尽管这些方法取得了进展，但训练的稳定性和计算规模限制其应用。

近年来，研究行人引入了全卷积神经网络U-Net，以增强对图像中语义信息的关注。U-Net在一定程度上提高了模型的准确性，但其单一的网络结构在处理过程中导致大量冗余计算，从而造成空间信息的丢失。为了解决U-Net的空间信息丢失问题，Cho等人[7]提出了一种从粗到精的策略，以更好地保留图像的特征信息。然而，这种方法在捕捉长距离依赖关系方面的效率和准确性仍有待提高。

随着 Transformer 的发展，自注意力机制也引入到了图像去模糊领域。由Zamir等人[8]提出的基于Transformer的高效模型Restormer，在高质量图像恢复任务中取得了有希望的结果。然而，传统的Transformer模型在处理像素级数据时，难以有效捕捉局部像素之间的关系，导致不必要的计算开销。

因此，提出了一种基于自注意力的非对称多尺度U-net。通过引入注意力机制，提高了整体图像去模糊的准确性，并通过多尺度架构更大程度地保留了空间信息。

2 Related Work

Multi-scale deep convolutional network deblurring methods

深度多尺度卷积神经网络，作为端到端学习策略的先驱性探索，成功地利用深度卷积网络处理不同尺度的图像。这种网络设计精妙地采用了一种由粗到精的策略，逐步恢复图像的清晰度。在这个过程中，网络的每一层的输入巧妙地将粗尺度网络的输出与细尺度网络的输入连接起来，实现了从粗到细的信息的平滑传递。网络每一层的处理严格遵循方程(1)。

通过上述改进，这两个网络为未来在多尺度处理方面的研究提供了新的思路和方向。

Transformer Deblurring Method

Transformer因其独特的全局注意力机制[10]而在捕捉信息中的长距离依赖方面表现出显著优势。与传统的卷积结构相比，Transformer能够访问更广阔的感受野，使模型能够关注上下文信息，从而生成更真实的图像输出。因此，研究行人将自注意力机制引入到图像去模糊任务中。

然而，简单地将自注意力机制应用于图像去模糊任务会导致计算成本显著增加。为了解决这个问题，引入了一种在特征深度域计算缩放点积注意力的方法，这使得能够从不同的特征通道高效提取特征信息。此外，还提出了结合多尺度分层模块的多头注意力机制，这不仅增强了模型的表示能力，而且有效地降低了计算成本。

尽管这些方法在精度方面取得了一些进展，但它们仍然需要经过复杂的矩阵乘法运算，其复杂性仍然是。因此，如何在保持性能的同时进一步降低计算复杂性仍然是未来研究的重要方向。可能的解决方案包括优化注意力机制的计算，采用更高效的特征提取方法，或引入新的模型架构。

Fourier Transform

卷积定理是信号处理中一个重要的原理，它表明时间域中的卷积操作与频率域中的乘法操作是等价的。这一特性同样适用于图像处理，使得时间域中复杂的卷积计算可以通过频域操作简化。具体来说，通过利用快速傅里叶变换（FFT）和逆快速傅里叶变换（IFFT），原本需要时间复杂度的卷积操作可以在的时间复杂度内完成。快速傅里叶变换的数学表达式见方程（5）。

f和g的卷积操作等同于执行快速傅里叶变换、，然后在频域中进行乘法操作，之后使用快速傅里叶逆变换将其转换回时间域。

在图像去模糊领域，许多研究已经充分利用了卷积定理的这一性质[12, 13, 14]。通过利用卷积定理和快速傅里叶变换，研究行人能够在保持性能的同时降低图像处理任务的计算复杂度。这种方法不仅提高了处理速度，还有助于构建更高效、轻量级的深度学习模型，为实际应用提供了更多可能性。

3 Proposed Method

图1：整体网络架构

在本文中，提出了一种基于自注意力的非对称多尺度U-net（AMSA-UNet），该方法结合了多输入多输出网络架构与 Transformer 模块，以解决由单尺度U-Net网络引起的图像空间特征丢失问题。该方法通过在解码器模块中引入自注意力来增强模型的感受野。同时，利用傅里叶变换提高模型的计算能力并降低其计算复杂度。此外，保留了传统多输入多输出-U-Net（MIMO-UNet）[15]架构中的多尺度特征融合模块，以进一步增强模型学习多尺度图像信息的能力和泛化能力。网络的整体框架如图1所示

编码器块

为了减少模型下采样过程中特征信息的丢失，并有效地解决不同尺度上的模糊问题，编码器采用了多尺度输入策略，将上层的下采样与上层输出作为当前层的输入。

在编码器块中，使用浅层卷积模块（SCM）从下采样的图像中进行特征提取，其中SCM是一个由两组3×3和1×1卷积层组合而成的通路，用于处理原始输入并将结果连接回原始输入后输出，SCM的结构如图2所示。SCM模块的输出和上层编码块的输出通过特征注意力模块（FAM）进行融合，这可以强调前一个尺度的有效信息或抑制不必要的信息，并从SCM的输出中学习空间信息，使得每一层的编码器块的输入包含最有效的多尺度信息。因此，FAM模块的输入将被作为这一层的编码器块的输入。具体的FAM结构如图2所示。

由于并非所有的频率信息都对去模糊任务有效，因此引入了一种基于频域的前馈神经网络（DFFN）作为编码器块的主要组成部分，以确定哪些频率信息应该保留以更好地重建锐利图像。为了区分有效的频率信息，受到压缩图像算法的启发，在DFFN中引入了一个可学习的量化矩阵W，它通过压缩方法的逆来学习确定哪些频率信息是有效的。在前馈通道中，输入通过FFT转换为频域信息，然后与量化矩阵W相乘。最终的输出是输入和通道输出的组合，DFFN模块的输出将作为当前层编码器块的最终输出的一部分，供下一层的多尺度输入使用。详细的DFFN模块结构如图2所示。

Decoder block

在ASMA-UNet中，解码器由三层解码块组成。每层解码块输出不同尺度结果，并将低层输出与上层特征图融合作为当前层解码块的输入，可以实现中间监督在各个解码块上的应用，提高模型输出结果的准确性。

每个层解码块的输入来自下层层解码块的输出以及通过融合提取（Fuse）模块的非对称特征融合（AFF）模块的输出。AFF模块通过1×1卷积层和3×3卷积层连接三层编码器模块的输出，以获得不同尺度的输出，从而使不同尺度的信息得到充分利用。在网络中，第一层和第二层解码块将接受相应尺度AFF模块的输出和下层解码模块的输出，由特征融合模块Fuse处理后作为该层解码模块的输入。在特征融合模块Fuse中，它将AFF模块的输出和下层解码模块的输出连接起来，输入到DFFN模块中，然后经过分割和求和得到融合后的最终结果。与传统的U-Net网络直接融合同层编码器的输出和解码器的输出不同，AFF模块的存在使得融合结果包含更多尺度上的空间信息，提高了解码块的准确性。Fuse和AFF的结构如图3所示。

解码块的主要组成部分是DFFN模块和基于频域的自注意力解算器（FSAS）。对于FSAS模块，传统的视觉 Transformer 通常通过将线性变换、和应用于输入特征X来计算、和，然后引入扩展函数提取和以及，并对它们进行 Reshape 操作以获得 Query Q、关键词K和值V，并通过缩放的点积注意力机制实现基于这三者的特征图的提取。计算过程如方程（6）所示。

图2：编码器块的结构。（a）SCM，（b）FAM，（c）DFFN。

其中R是 Reshape 函数，和分别表示捕获的 Patch 块的高度和宽度。这种方法计算复杂度高，计算量大，一旦捕获的 Patch 数量增多，其局限性就会变得明显。因此，在FASA模块中，通过一组1×1卷积和3×3深度卷积获得Q、K和V，Q和K通过FFT逐元素相乘，替换方程（4）中的，大大减少了其计算量。然后将结果反转并相乘以获得估计的特征图。这通过频域中的逐元素乘法操作实现了注意力图的快速估计，而不是空间域中的矩阵乘法。详细的FSAS结构如图3所示。

Asymmetric U-shaped network architecture

不对称的U-net结构体现在编码器模块和解码器模块之间的不对称性。如前一小节所述，解码器模块包含DFFN模块和FASA模块，而编码器模块仅使用DFFN模块。这是因为Transformer模块更适合解码器进行长距离依赖捕捉，从而生成更高质量的结果。而编码器模块提取更浅的特征，这些特征通常包含模糊效果，如果将FSAS应用于编码器模块，它会将清晰特征与模糊特征混淆，这对于图像处理是不利的。因此，网络采用不对称结构以实现更好的去模糊效果。编码器部分遵循方程(7)，解码器部分遵循方程(8)。

图3：解码器块的结构。(a) AFF，(b) Fuse，(c) FSAS。

4 Experiments

Dataset and implementation details

在本论文中，作者使用了GoPro[4]训练数据集进行训练，该数据集包含2103对模糊清晰的控制图像。用于训练的验证集是460对随机划分的图像。测试集包括来自GoPro的1111对模糊清晰图像和来自Kohler[19]测试数据集的48对模糊清晰图像。在每次训练迭代中，经过随机裁剪操作后，随机选择4张图像作为输入，以获得256×256大小的图像。初始学习率设置为，并且每500个周期降低0.5倍。模型在GoPro训练集上训练了1200次迭代以得到最终结果。除了将GoPro训练得到的模型应用于划定的GoPro测试集之外，同样的模型也直接应用于Kohler测试集，以此作为检查模型泛化能力的一种方式。

Performance comparison

ASMA-UNet进行比较。在GoPro数据集上的测试结果如表1所示，本文提出的ASMA-UNet在其余去模糊方法中表现出色。通过引入自我关注机制，提高了准确性，模型在捕捉长距离依赖关系方面取得了巨大进步。ASMA-UNet在测试集上取得了平均PSNR 30.55dB的成绩，比基于三种对抗网络架构的模型（即FCL-GAN、DeblurGAN和DeblurGANv2）的平均PSNR值分别高出5.72dB、1.33dB和1.01dB，平均SSIM与DeblurGAN和FCL-GAN相比也有显著提高，分别提高了0.17和0.08。与基于单网络层结构的SVRNN、SRN和SVDN相比，所提方法分别实现了平均PSNR改进1.37 dB、0.30 dB和0.75 dB。

这三种方法得益于单一网络架构，在处理速度上优于其他方法。例如，与DeepDeblur这样的深度卷积方法相比，平均处理速度提高了近4秒，SVDN的平均运行时间仅为0.01秒，比DeepDeblur快4.32秒。然而，得益于引入的频域计算方法，作者的方法的平均运行时间为0.05秒，仅比SVDN慢0.04秒，同时比SRN和SVRNN分别快1.83秒和1.35秒。基于多尺度处理方法的PSS-NSC方法，平均PSNR略高0.36 dB，但平均处理时间比提出的方法多近1秒。总之，ASMA-UNet在保证准确性的同时，实现了减少运行时间的目标。图4展示了来自三个经典网络--GAN、单尺度卷积和多尺度卷积--的代表性方法的比较结果。

图4：GoPro上的可视化比较。为了清晰起见，展示了结果图像的放大部分。(a) 模糊，(b) DeblurGAN，(c) PSS-NSC，(d) Deepdeblur，(e) 本文提出的方法，(f) 真实情况。

可以看出，与DeblurGAN、DeepDeblur和PSS-NSC方法相比，本文提出的方法在视觉方面表现更好。从结果来看，本文方法的处理效果明显比DeBlurGAN和DeepDeblur少伪影，边缘更清晰，基本上恢复了图像轮廓及相关信息。在恢复远处视图中文本信息的性能方面，与PSS-NSC相比，本文方法具有更好的处理效果。在GoPro数据集上训练的ASMA-UNet和DeblurGAN直接在Kohler数据集上进行测试。

结果表明，与DeblurGAN相比，ASMA-UNet产生了更清晰的边缘，更少的伪影和波纹噪声。这表明本文提出的方法具有更强大和更稳定的泛化能力。视觉结果如图5所示。

Ablation Experiment

为了彻底分析本文提出的非对称多尺度特征融合（AFF）和非对称架构（AA）的有效性，进行了消融实验。在这些实验中，从GoPro训练集中随机选择了一个减少的数据集，包含480对模糊和清晰图像，以及一个包含200对图像的测试集。在减少的数据集上对每个组件评估了150次迭代，比较结果如表2所示。首先，通过在编码器中添加与解码器中相同的DFFN模块，测试了对称架构，形成了对称的网络结构。其次，移除了AFF模块，使得每个解码层的输入仅由前一层输出的上采样形成。结果如下表2所示。

可以观察到，当网络架构对称时，与原始网络模型相比，PSNR下降了0.09 dB，运行时间增加了0.01秒。这直接与在编码器模块中引入DFFN有关，这不仅导致实验结果下降，还导致运行时间增加。当移除AFF时，与原始模型相比，PSNR下降了0.26 dB，确认了多尺度融合模块在网络架构中对处理效果的重要性。视觉结果如图6所示，证明了原始模型取得了最佳的视觉效果。

5 Conclusion

如上所述，提出了一种基于自注意力的非对称多尺度U-net，该网络融入了自注意力机制，用于去模糊任务。

与仅基于卷积神经网络的方法或直接集成自注意力机制的方法相比，AMSA U-Net在准确性和速度之间取得了优越的平衡，可以在扩大模型感受野的同时提高计算效率，从而实现有效的去模糊。

解码器将多尺度特征融合结果作为输入，通过融合自注意力机制的前馈网络进行处理，并产生输出。

这使得网络能够更好地捕捉不同尺度上的特征信息。比较结果表明，所提出的方法显著优于其他八种优秀方法。

参考

[1].AMSA-UNet: An Asymmetric Multiple Scales U-net Based on Self-attention for Deblurring.