LKDN：大型内核蒸馏网络，实现高效的单图像超分辨率

LKDN：大型内核蒸馏网络，实现高效的单图像超分辨率

Large Kernel Distillation Network for Efficient Single Image Super-Resolution

大型内核蒸馏网络，实现高效的单图像超分辨率

https://arxiv.org/abs/2407.14340v1#

https://github.com/stella-von/LKDN

CVPRW 2023

摘要

• 问题：目前最先进的（SOTA）模型仍然面临计算成本高等问题
• 方法：本文提出了大型内核蒸馏网络（Large Kernel Distillation Network LKDN），
  • 简化模型结构，引入更高效的注意力模块，称为大内核注意力（large kernel attention LKA），以降低计算成本，同时提高性能
  • 采用重参数化技术来提高模型性能，而不会增加额外的成本
  • 引入新的优化器 Adan，提高训练速度和性能
• 结果：实验结果表明，LKDN的性能优于现有的轻量级SR方法，并实现了SOTA性能。
• 主要贡献：
  • 在分析 BSRN 和 VAPSR 的计算效率后，我们通过简化模型结构和引入更高效的注意力模块，在减少参数数量和计算消耗的同时实现更好的性能
  • 使用重参数化技术来提高模型性能，而不会引入额外的推理负担
  • 引入新的优化器，可以同时提高 SISR 模型的训练速度和性能

1. 引言

图1：Urban100 与 SR X4 的模型性能和复杂度的比较

在轻量级超分辨率模型的众多设计方法中，信息蒸馏连接已被确定为一种非常有效的方法。这种方法融合了不同层次结构的特征，以促进将更多独特的特征传输到网络中，从而实现高效的特征重用，并在重建精度和计算效率之间实现更好的平衡。

2. 相关工作

2.1 高效的 SR 模型

轻量级SR网络的常见策略：

• 网络剪枝
• 参数共享
• 知识蒸馏
• 深度可分离卷积、蓝图卷积
• 注意力机制
• 高效的上采样方法
• 重参数化技术

2.2 大型内核设计

2.3 重参数化

2.4 Adan 优化器

Adam 优化器在各种深度学习领域中得到了广泛的应用。通过利用一阶和二阶梯度距估计，动态调整每个参数的学习率，以实现比随机梯度下降更快的收敛。然而，Adam 也可能受到非收敛和局部最优的影响。最近，Adan 优化器结合了修正的 Nererov 脉冲、自适应优化和解耦权重衰减。Adan 使用extrapolation points 预先感知梯度信息，从而可以有效地逃避尖锐的局部最小值并提高模型泛化度。基于大量的实验，已经表明 Adan 优化器在 CNN 和 Transformer 上都优于现有的 SOTA 优化器。因此，我们的目标是将 Adan 算法应用于轻量级超分辨率任务。

3. 方法

3.1 网络架构

所提出的方法遵循与 BSRN 相同的框架设计，包括 浅层特征提取、多堆叠特征蒸馏块、多层特征融合和图像重建块，四个部分：

1. 浅层特征提取

   \[I^n_{LR} = Concat(I^n_{LR}), \\ F_0 = h_{ext}(I^n_{LR}) \]

2. 多堆叠特征蒸馏块

   \[F_k = H^m_{LKDB}(\ldots H^1_{LKDB}(F_0)), 1\le k\le m \]

3. 多层特征融合

   \[F_{fusion} = H_{fusion}(Concat(F_1,\ldots, F_k)) \]

4. 图像重建块

   \[I_{SR} = H_{rec}(F_{fusion} + F_0) \]

3.2 重新思考 BSRN

通过ESA、CCA和多残差连接，BSRN模型的性能得到了提高。然而，结构复杂导致计算效率较低。RLFN使用基于一次性结构化剪枝算法的剪枝敏感性分析工具分析了ESA块的冗余度，并发现了大量的冗余度。因此，我们去掉了BSRN的ESA和CCA模块，引入了更高效的注意力模块。

最近的研究表明，通过合理使用大核卷积，可以提高模型的性能，同时保持接收的复杂度。VAN利用卷积分解将大型内核卷积分为三个部分：空间局部卷积、空间远程卷积和通道卷积。具体来说，\(K\times K\) 卷积被分解成 \((2d-1)\times (2d-1)\) 深度卷积、膨胀系数为 d 的 \([\frac{K}{d}]\times [\frac{K}{d}]\) 深度膨胀卷积和 \(1\times 1\) 卷积。通过分解大型内核卷积，该模型可以以最小的计算成本和参数捕获长距离关系。与 VAPSR 类似，我们以不同的顺序进行卷积分解，因此，LKA 模块可以表示如下：

\[X_{atten} = Conv_{DW-D}(Conv_{DW}(Conv_{1\times 1}(F))) , \]

\[Output = X_{atten} \otimes F \]

• 通过将 \(17\times 17\) 大型卷积分解为 \(1\times 1\) 卷积、\(5\times 5\) 深度卷积和大小为 5 膨胀系数为 3 的深度膨胀卷积

3.2.1 重参数化

为了替换 BSRN 中的蓝图浅层残差块（BSRB），我们引入了重参数化的蓝图浅层块（RBSB）结构。具体来说，我们在 \(1\times 1\) 逐点卷积和 \(3\times 3\) 深度卷积上引入了一个可重新参数化的 skip 连接。此外，我们在 \(3\times 3\) 深度卷积上并行了一个 \(1\times 1\) 深度卷积。

重新参数化包括两个步骤

1. 输入特征 \(F_0\) 通过重参数化 \(1\times 1\) 逐点卷积传递

   \[F_1 = Conv_{1\times 1}(F_0) + F_0 \]

2. 应用重参数化 \(3\times 3\) 深度卷积

   \[F_2 = Conv_{DW_{3\times 3}}(F_1) + Conv_{D_{1\times 1}}(F_1) + F_1 \]

3.2.2 大型内核蒸馏块

在分析了 BSRN 和 VAPSR 的特性后，我们开发了一种更高效的大型内核蒸馏块 （LKDB）。它包括四个部分：特征蒸馏、特征融合、特征增强和特征转换。

1. 特征蒸馏：

2. 特征融合

3. 特征增强

4. 特征转换

• 为了提高模型的性能，在特征转换阶段采用 1x1 卷积，同时加入像素归一化模块，以确保稳定的模型训练

5. 长跳跃连接来加强模型的残差学习能力

4. 实验

4.1 数据集和评估指标

• 数据集：DF2K
• 数据增强：随机水平翻转和90度旋转

4.2 LKDN 实现细节

LKDN：使用 BSB 进行训练以减少训练时间

• blocks ：8

• channel_num : 56

• attention_num : 56

• batch_size : 64

• patch_size : 48x48

• L1 loss function

• Adan 优化器

  • \(\beta_1=0.98,\beta_2=0.92,\beta_3=0.99\)
  • EMA = 0.999

• iterations ：1e6

• learning_rate ：5e-3

LKDN-S：使用RBSB进行训练

• blocks ：5

• channel_num : 42

• Adan 优化器

  • \(\beta_1=0.98,\beta_2=0.92,\beta_3=0.99\)
  • EMA = 0.999

• 预训练阶段：

  • patch_size : 64x64
  • L1 loss function
  • learning_rate : 5e-3
  • iterations : 9.5e5

• 微调阶段:

  • patch_size : 120x120

  • batch_size : 64

  • L2 loss function

  • iterations : 5e4

4.3 消融实验

4.3.1 大型内核注意力

• C：蒸馏结构的输入通道数
• A：注意力模块的输入通道数
• 结论：与原始 BSRN 相比，LKDN 可以在 Urban100 和 Managa109 上实现超过 0.1 dB的性能增益，并保持较低的参数和计算成本。

4.3.2 重参数化

• 结论：消除不必要的残余连接可以提高性能，同时降低模型复杂性。因此，重参数化可以进一步提高性能，同时保持模型复杂性。

4.3.3 优化

• 结论：使用 Adan 优化器的训练速度提高了 10%，在各种基准数据集上的性能都有了显著提高。此外，Adan优化器可以实现更快的收敛

4.4 与最先进的方法的比较

• 定量比较

• 定性比较

• 推理速度比较

5. 结论

我们提出了大核蒸馏网络（LKDN）作为一种高效的单像超分辨率（SISR）解决方案。通过对一些最先进的轻量级模型的仔细分析，我们发现了它们的弱点，并对其进行了改进，以提高LKDN的性能。通过结合大核设计，简化模型结构，引入更高效的注意力模块，并采用重新参数化，我们在性能和计算效率之间实现了平衡。此外，还引入了一种新的优化器来加速LKDN的收敛。大量实验表明，LKDN 在性能、参数和 Multi-Adds 方面优于最先进的高效 SR 方法。