RFDN：用于轻量级图像超分辨率的残差特征蒸馏网络

RFDN：用于轻量级图像超分辨率的残差特征蒸馏网络

Residual Feature Distillation Network for Lightweight Image Super-Resolution

轻量级图像超分辨率的残差特征提取网络

https://arxiv.org/abs/2009.11551

https://github.com/njulj/RFDN

ECCV 2020

Abstract

• 问题：基于 CNN 的方法取得了巨大成功，但由于需要大量计算，将这些方法应用于边缘设备并不容易
• 方法：重新思考信息多重蒸馏网络 （IMDN），并提出一种轻量级且准确的 SISR 模型，称为残差特征蒸馏网络 （RFDN）
  • 提出了功能蒸馏连接 （FDC），它在功能上等同于通道拆分作，同时更加轻量级和灵活
  • RFDN 使用多个特征蒸馏连接来学习更多判别性特征表示
  • 提出了一个浅层残差块 （SRB） 作为 RFDN 的主要构建块
• 结果：所提出的 RFDN 在性能和模型复杂度方面与最先进的方法相比实现了更好的权衡
• 主要贡献：

  • 我们提出了一种 轻量级残差特征蒸馏网络（RFDN），以实现快速准确的图像超分辨率，该网络实现了最先进的SR性能，同时使用的参数比竞争对手少。

  • 对 信息蒸馏机制（IDM） 进行更全面的分析，对IMDN网络进行重新思考。基于这些新的理解，我们提出了比IDM更轻量级和更灵活的 特征蒸馏连接（FDC）。

  • 我们提出了 浅残差块（SRB） ，它将恒等式连接与一个卷积块结合在一起，以进一步提高SR性能，而无需引入任何额外的参数。

1. 引言

信息蒸馏网络：模型中增强单元（enhandcement unit）架构。橙色圆圈表示切片操作，紫色圆圈表示通道维度上的串联操作。

为了减少参数数量，DRCN（用于图像超分辨率的深度递归卷积网络）和 DRRN（通过深度递归残差网络实现图像超分辨率），通过参数共享有效地减少参数数量。但是，它必须增加网络的深度或宽度，以补偿递归模块造成的损失。这些模型减少了模型大小，但代价是增加了操作数量和推理时间。

Hui等提出了一种信息蒸馏网络（information distillation network IDN），该网络沿通道维度将中间特征明确地分为两部分，一部分被保留，另一部分通过后续卷积层进行进一步处理。通过使用这种通道拆分策略，IDN 可以将当前信息与部分保留的本地短路径信息聚合在一起，并在适度的尺寸下实现良好的性能。后来，IMDN通过设计一个信息多蒸馏块 （information multi-distillation block IMDB） 来进一步改进 IDN，该模块在粒度级别提取特征。具体来说，通道拆分策略在IMDB中被多次应用。每次，都会保留一部分功能，并将另一部分发送到下一步。

IDN （通过信息蒸馏网络实现快速准确的单图像超分辨率）和 IMDN（具有信息多蒸馏网络的轻量级图像超分辨率） 的关键组成部分是信息蒸馏机制 （information distillation mechanism IDM），它将前面提取的特征明确地分为两部分，一部分保留，另一部分进一步完善。我们认为IDM不够高效，并且在网络设计中带来了一些不灵活性。很难将特征连接与 IDM 合并。本文将对信息蒸馏机理进行更全面的分析，并提出比IDM更轻量级、更灵活的特征蒸馏连接（FDC）。我们使用 IMDN 作为基线模型，因为它在重建质量和推理速度之间做出了很好的权衡，非常适合移动设备。但IMDN还不够轻量级，SR性能还有待进一步提升。为了构建更强大、更快速、更轻量级的 SR 模型，我们重新思考了 IMDN 的架构，并提出了残差特征蒸馏网络 （residual feature distillation network RFDN）。与 IMDN 相比，我们的 RFDN 通过使用特征蒸馏连接 （feature distillation network FDC） 更加轻量级。此外，我们提出了一种浅残差块（SRB），它用作RFDN的构建块，以进一步提高SR性能。SRB 由一个卷积层、一个相同的连接和末端的激活单元组成。与普通卷积相比，它可以从残差学习中受益，而无需引入额外的参数。将 SRB 与功能蒸馏连接相结合非常容易，以构建更强大的 SR 网络。

图1：残差特征蒸馏网络（RFDN）的架构

2. 相关工作

3. 方法

3.1 信息多蒸馏块（information multi-distillation block）

图2 （a）IMDB ：原始信息多蒸馏块 the original information multi-distillation

​ （b）IMDB-R ：对IMDB的重新思考 rethinking of the IMDB

​ （c）RFDB ：残余特征蒸馏块 residual feature distillation block

​ （d）SRB ：浅层残余块 shallow residual block

如图2（a）所示，信息蒸馏块（IMDB）的主要部分是渐进式细化模块（progressive refinement module），其标有灰色背景。PRM首先使用 \(3\times 3\) 卷积层来提取多个后续蒸馏步骤的输入特征。对于每个步骤，对上述特征采用通道分离操作，并将输入特征分为两部分。一部分被保留，另一部分被送入下一个蒸馏步骤。给定输入特征 \(F_{in}\)，此过程可以描述为：

\[ F_{distilled\_1},F_{coarse\_1} = Split_1(L_1(F_{in}))\\F_{distilled\_2},F_{coarse\_2} = Split_2(L_2(F_{coarse\_1})) \\F_{distilled\_3},F_{coarse\_3} = Split_3(L_3(F_{coarse\_2})) \\F_{distilled\_4} = L_4(F_{coarse\_3}) \]

• \(L_j\) ：第 \(j^{th}\) 个卷积层（包括激活单元）
• \(Split_j\) ：第 \(j^{th}\) 个通道分割操作
• \(F_{distilled\_j}\) ：第 \(j^{th}\) 个特征蒸馏
• \(F_{coarse\_j}\) ：后续层将进一步处理的第 \(j^{th}\) 个粗略特征（coarse feature）

最后，所有蒸馏的特征被连接在一起，作为PRM的输出

\[F_{distilled} = Concat(F_{distilled\_1},F_{distilled\_2},F_{distilled\_3},F_{distilled\_4}) \]

• Concat ：表示沿通道维度的串联运算

3.2 重新思考IMDB

尽管 PRM 取得了显著的改进，但由于通道分离操作（channel splitting operation）不够高效和引入了一些不灵活性。蒸馏特征由具有许多冗余参数（redundant parameters）的 \(3\times 3\) 卷积过滤器生成（将\(3\times 3\)卷积改为\(1\times 1\)卷积）。此外，特征优化管道（feature refinement pipeline）（沿着PRM的右侧分支）与 通道分离操作 耦合在一起，因此很难仅对该管道使用 特征连接（identity connections）。接下来，我们重新思考通道分离操作，并给出一个新的等效PRM架构来解决上述问题。

如图2（b）所示，\(3\times 3\) 卷积层之后跟着通道分离层，然后可以解耦为两个 \(3\times 3\) 卷积层 \(DL\) 和 \(RL\) ，\(DL\) 层负责产生蒸馏特征，\(RL\) 层进一步处理后续粗略特征的细化层。整个结构可以描述为：

\[&& F_{distilled\_1},F_{coarse\_1} = DL_1(F_{in}),RL_1(F_{in}) \\ && F_{distilled\_2},F_{coarse\_2} = DL_2(F_{coarse\_1}),RL_2(F_{coarse\_1}) \\ && F_{distilled\_3},F_{coarse\_3} = DL_3(F_{coarse\_2}),RL_3(F_{coarse\_2}) \\ && F_{distilled\_4} = DL_4(F_{coarse\_3}) \]

上面的方程式描述了每组拆分操作可以看作是两个同时工作的卷积层。我们称这种新架构为IMDB-R，它比原来的IMDB更灵活。它对 PRM 的工作原理有了更清晰的了解，因此我们可以获得更多关于如何设计更高效的 SR 模型的线索。

3.3 残余特征蒸馏块（residual feature distillation block RFDB）

受IMDB重新思考的启发，在本节中，我们介绍了比IMDB更轻量级和更强大的残余特征蒸馏块（RFDB）。在图2中，我们可以看到信息蒸馏操作实际上是通过 \(3\times 3\) 卷积实现的，该卷积以固定比例压缩特征通道。然而我们发现使用 \(1\times 1\) 卷积进行通道缩减更有效，就像在许多其他CNN模型中所做的那样。如图2（c）所示，左边的三个卷积被 \(1\times 1\) 卷积取代，这大大减少了参数的数量。最右边的卷积仍然使用 \(3\times 3\) 卷积核，这是因为它位于 RFDB 的主体上，并且必须考虑空间上下文以更好地优化特征。我们将这些外部连接称为 特征蒸馏连接（feature distillation connections FDC）。

尽管有上述改进，我们还在网络中引入了更细粒度的残差学习。为此，我们设计了一个 浅残差块（shallow residual block SRB），如图2（d）所示，它由 \(3\times 3\) 卷积、特征连接（identity connectin）和 激活单元 组成。SRB可以从残差学习中受益，而无需引入任何额外的参数。原始IMDB仅包含中级残差连接（mid-level residual connection），这些连接对于网络来说太粗糙了，无法从残差连接中受益最多。相比之下，我们的SRB支持更深的残差连接，即使用轻量级的SR模型（lightweight shallow SR model），也可以更好地利用残差学习的力量。

3.4 框架

图1：残差特征蒸馏网络（RFDN）的架构

我们使用与IMDN相同的框架，如图1所示，残差特征蒸馏网络（RFDN）由4部分组成：特征提取卷积（feature extraction convolution），多堆叠残差特征蒸馏块（multiple stacked residual feature distillation blocks RFDBs），特征融合部分（feature fusion part），重构块（reconstruction block）。

• 特征提取模块

  • 通过 \(3\times 3\) 卷积实现初始特征提取，以从输入的LR图像生成粗略的特征。

  • 给定输入 \(x\) ，此过程可以表示为

    \[F_0 = h(x) \]
    • \(h\) ：粗略特征提取函数
    • \(F_0\) ：提取的特征

• 多堆叠残差特征蒸馏块

  • 多个RFDB以链式方式堆叠以逐渐细化提取的特征。

  • 这个过程可以表述为

    \[F_k = H_k(F_{k-1}), K=1,\cdots,n \]
    • \(H_k\) : 表示第 \(k^{th}\) 个 RFDB 函数
    • \(F_{k-1}\) : 表示第 \(k^{th}\) 个RFDB的输入特征
    • \(F_{k}\) : 表示第 \(k^{th}\) 个RFDB的输出特征

• 特征融合模块

  • 经过RFDB的逐步细化后，所有中间特征都由 \(1\times 1\) 卷积层组装。然后，使用 \(3\times 3\) 卷积层对聚合特征进行平滑处理。

  • 如下所示

    \[F_{assemble} = H_{assemble}(Concat(F_1,\cdots,F_n)) \]
    • \(Concat\) : 是沿通道维度的串联运算
    • \(H_{assemble}\) : 表示 \(1\times 1\) 卷积后跟 \(3\times 3\) 卷积
    • \(F_{assemble}\) : 聚合特征（aggregated features）

• 重构块

  • 通过重建生成SR图像

  • 如下所示：

    \[y = R(F_{assemble} + F_0) \]
    • \(R\) : 表示重建函数
    • \(y\) : 表示网络输出
    • 重建过程仅由 \(3\times 3\) 卷积和 亚像素卷积 组成

• 损失函数

  \[\mathbb{L}(\theta) = \frac{1}{N}\sum^{N}_{i=1}|| H_{RFDN}(I^{LR}_{i}) - I^{HR}_{i} ||_1 \]
  • \(H_{RFDN}\) : 网络
  • \(\theta\) : RFDN的可学习参数
  • \(|| \ ||_1\) : \(l_1\) 范数
  • \(I^{LR}\) : 输入的LR图像
  • \(I^{HR}\) : 对应的真实有效的HR图像

4. 实验

4.1 数据集和指标

• 数据集：DIV2K
• 测试集：Set5，Set14， BSD100，Urban100，Managa109
• 指标：PSNR峰值信噪比 SSIM结构相似性
  • 所有的值都是在RGB通道转换而来的YCbCr的Y通道上计算的

4.2 实现细节

• 数据集：

  • 随机裁剪64个64x64的补丁作为每个训练小批量的输入
  • 随机水平翻转和90°旋转来增强训练数据集

• 优化器

  • Adam
  • \(\beta_1=0.9 ,\ \beta_2=0.999 ,\ \epsilon=10^{-8}\)
  • 学习率初始化为 \(5\times 10^{-4}\)
  • 每 \(2\times 10^5\) 次小批量更新时减半

• 预训练

  • \(\times 2\) 模型是从头训练的
  • 模型收敛后，将其用作其他尺度的预训练网络
  • 消融研究中的所有模型都是从头开始训练的，以节省训练时间

• RFDN

  • 通道数：48
  • RFDB：6

• RFDN-L

  • 通道数：52
  • RFDB：6

4.3 模型分析

4.3.1 消融研究

图3：消融实验中使用的 Base块，SRB块，FDC块 和 RFDB。

为了评估所提出的特征蒸馏连接（FDC）和浅残差块（SRB）的重要性，我们设计了四个块，它们将分别堆叠为SR网络的主体部分（图1）。图 3 描绘了四个模块，评估结果如表 1 所示。比较表 1 的前两行，我们可以发现 SRB 在不引入任何额外参数的情况下提高了性能（例如 Manga109的 PSNR：+0.12dB， SSIM：+0.0024）。在比较最后两行时，我们也可以观察到类似的改进，这表明浅残差块的有效性。通过添加FDC，基本方法的性能得到了很大的提高，例如Manga109的PSNR从30.28提高到30.47（+0.19dB）。得益于 FDC 和 SRB，我们的 RFDB 性能明显优于基础模块。

表1：在比例因子为 x4 的基准数据集上对FDC和SRB的调查。突出显示最佳结果

4.3.2 蒸馏比例的调查

我们研究了表2中特征蒸馏连接的蒸馏比例。不同的蒸馏比例表示特征蒸馏连接中不同的输出通道数量。如表2最后一行所示，当蒸馏速率增加时，PSNR和SSIM在每个数据集上的生长趋势不同。总体而言，0.5 的蒸馏比率在 SR 性能和参数数量之间具有很好的权衡，在我们的 RFDN 和 RFDN-L 中采用参数0.5作为最终蒸馏比率。

表2：对比例因子为 X4 的基准数据集的蒸馏比例的调查，突出显示最佳结果。\(\uparrow\) 代表上升， \(\wedge\) 代表下降，\(\uparrow\ \wedge\) 代表先上升后下降。

4.3.3 模型复杂度分析

图 4 描述了 PSNR 与 Set5 × 4 数据集上参数的比较。图 4 所示的模型包括 DRRN、LapSRN、VDSR、MemNet、IDN、CARN-M和 IMDN。在评估轻量级模型时，模型参数的数量是一个需要考虑的关键因素。从表 3 中可以看出，与参数较少的最先进的轻量级模型相比，我们的 RFDN 实现了相当或更好的性能。如图 4 所示，尽管 IMDN 与之前的方法（如 MemNet 和 IDN）相比取得了显着的改进，但它的参数比大多数轻量级模型更多。相比之下，我们的 RFDN 在参数较少的情况下实现了比 VDSR、MemNet、IDN 和 IMDN 更好的性能。当使用更多的特征通道时，我们的 RFDN-L 在保持适度的模型尺寸的同时，取得了比 RFDN 更好的结果。为了更全面地了解模型的复杂性，我们还在图 5 中显示了 Set5 × 4 数据集上 PSNR 与 Mult-Adds 的比较。正如我们所看到的，我们的 RFDN 和 RFDN-L 实现了比 IMDN 更高的 PSNR，同时使用更少的计算。

4.4 与最先进的技术比较

表3：数据集 Set5、Set14、BSD100、Urban100 和 Manga109 上比例因子 2、3 和 4 的平均 PSNR/SSIM。最佳和次优结果分别以红色和蓝色突出显示。

4.5 关于实验设置

如第 4.2 节所述，我们在训练模型时使用略有不同的实验设置。为了更清楚地了解RFDN的改进，我们在相同的实验设置下从头开始训练RFDN和IMDN。表 4 显示了五个基准数据集的性能比较。我们的RFDN在所有数据集上都优于IMDN，无论是PSNR还是SSIM，参数都少得多，这证明我们的RFDN网络设计的改进确实提高了图像SR的性能。

表4：RFDN和IMDN在相同实验设置下的性能比较。这两个模型都是使用上采样系数 \(\times 4\) 从头训练的。

4.6 针对AIM20挑战的增强型RFDN

如表5所示，我们的增强型RFDN（E-RFDN）在AIM 2020高效超分辨率挑战赛中获得第一名。具体来说，我们用ESA块替换了RFDB中的CCA层，并在E-RFDN中使用了4个这样的增强型RFDB（E-RFDB）。E-RFDN 中的特征通道数设置为 50，特征蒸馏速率为 0.5。在 E-RFDN 训练期间，从 HR 图像中随机裁剪大小 256 × 256 的 HR 补丁，并将小批量大小设置为 64。E-RFDN 模型是通过使用 Adam 优化器最小化 L1 损失函数来训练的。初始学习率设置为 \(5\times 10^{-4}\) ，每 200 个周期减半。在 1000 个 epoch 之后，L2 损失用于微调，学习率为 \(1\times 10^{-5}\) 。DIV2K 和 Flickr2K 数据集用于训练 E-RFDN 模型。我们在表5中包括了前五种方法，“#Activations”测量卷积层所有输出的元素数。与 AIM 2019 约束 SR 挑战中的第一名方法 IMDN 相比，我们的方法在运行时、参数、FLOP 和激活方面提供了显着的增益。更多细节和建议可以在中找到。

表5：AIM 2020 高效SR挑战结果

5. 结论

本文对轻量级图像超分辨率的信息蒸馏机理进行了全面分析。然后，我们重新考虑了信息多蒸馏网络（IMDN），并提出了更加轻量级和灵活的特征蒸馏连接（FDC）。为了进一步提高超分辨率性能，我们还提出了浅残差块（SRB），它将恒等式连接与一个卷积块相结合。通过使用浅层残差块和特征蒸馏连接，我们构建了残余特征蒸馏网络（RFDN），以实现快速轻量级的图像超分辨率。大量的实验表明，所提出的方法在定量和定性方面都取得了最先进的结果。