CVPR2021|超分性能不变，计算量降低50%，董超等人提出用于low-level加速的ClassSR

最新推荐文章于 2023-05-31 11:18:00 发布

AIWalker-Happy

最新推荐文章于 2023-05-31 11:18:00 发布

阅读量1.1k

点赞数 1

分类专栏：深度学习文章标签：深度学习 pytorch 图像处理

本文链接：https://blog.csdn.net/huohu728/article/details/115275312

版权

深度学习专栏收录该内容

30 篇文章 6 订阅

订阅专栏

编辑：Happy
首发：AIWalker

标题&作者团队

本文是深圳先进技术研究院孔祥涛、赵恒源、乔宇、董超等人在low-level领域关于超分网络加速的一次探索。它创新性的将分类与超分进行了融合，根据不同子块的复原难度自适应选择合适的超分分支以降低整体计算复杂度：复原难度低的平坦区域选择复杂度低的超分分支，复原难度高的纹理区域选择复杂度高的超分分支。在不降低超分性能的情况下，该方法可以最高可以节省50%的计算量。值得各位low-level领域的同学仔细研究一番。

Abstract

本文旨在加速超分模型在大图像(2K-8K)上的推理速度。在实际应用中，大图像同尝会被拆分为多个小图进行处理。基于该处理过程，我们发现：不同图像区域具有不同的复原难度，可以被不同容量的模型进行处理。直觉上来讲，平滑区域要比复杂纹理区域更易于超分。

为利用该特性，我们采用合适的SR模型处理拆分后的不同子图，即平滑区域采用简单的小模型，纹理区域采用复杂的大模型。以此为根基，我们提出了一种新的解决方案：ClassSR，它将分类与超分组合到统一框架中。具体来讲，它先采用ClassModule将子图按照复原难度分到不同的类别；然后对不同类别的子图采用不同的SRModule进行超分。这里的ClassModule就是常规分类网络，而SRModule则是由不同大小的SR模型构成。

我们更进一步引入一个带Class-Loss与Average-Loss的分类方法以生成分类结果。通过联合训练，大部分子图将被送入更小的网络。我们通过证实：在DIV8K数据集上，所提ClassSR可以帮助大部分现有方法(如FSRCNN、CARN、SRResNet、RCAN)节省50%FLOPs。注：这是一种通用性方法，同样适用于其他low-level视觉任务。

上图给出了本文所提方法处理示意图以及其对于不同超分方法的指标-计算量的调整示意图。本文的主要贡献包含以下三点：

提出ClassSR，首个将分类与超分结合的图像超分方案；
提出结合数据特性克服low-level加速问题，而且它与其他加速计算具有互补性；
提出一种带两种新颖损失的分类方案以促进子图按照复原难度进行分类。

Method

我们首先来看一下DIV2K中不同子图(将DIV2K拆分为 $32\times 32$ 大小的子图)的类型，见上图与表。为评估不同子图的复原难度，我们采用与训练MSResNet进行评估，按照其PSNR值进行排序。如图所示，我们将其均匀的分为三组：simple、medium以及hard。很明显：具有更高PSNR的子图通常是平坦区域，而具有更低PSNR的子图包含复杂的纹理。

我们采用不同的网络处理不同类型的子图，如上表所示。我们采用三个不同大小的FSRCNN(通过通道数控制大小)进行处理。这些模型分别在不同类型的子图上进行训练。可以看到：在simple子图上，FSRCNN(16)与FSRCNN(56)的性能基本没区别；类似的，FSRCNN(36)可以在medium子图上取得了与FSRCNN(56)相当的性能。这也就意味着：我们可以采用轻量型网络处理simple子图以节省计算量。这也就是本文所提ClassSR的初衷与核心所在。

ClassSR

上图给出了本文所提ClassSRC的网络架构示意图，它包含两个关键性模块：ClassModule与SRModule。ClassModule负责将输入子图像划分到M个类别，而SRModule则包含M个分支(即不同复杂度的SR网络) ${f_{SR}^j\}_{j=1}^M$ 。具体来讲，大的输入图像LR图像X首先被拆分为重叠子图像 ${x_i\}_{i=1}^N$ ；ClassModule对接收到的每个子图进行分类，输出概率向量 $[P_1(x_i), \cdots, P_M(x_i)]$ ；我们选择概率最大(即J=argmax_j P_j(x_i))的那个分支SR网络对当前块进行超分 $y_i = f_{SR}^J(x_i)$ ；最后我们组合所有超分子图得到最终的大的SR图像 $Y$ 。

ClassModule

该模块的作用是告诉我们“输入子图是比较容易重建还是难以重建”。我们设计了一个简单的分类网络，包含5个卷积层、一个均值池化以及全连接层。卷积层负责提取特征，池化与全连接层负责输出概率向量。该模块极为轻量，故而仅带来很少的计算量。实验表明：如此简单的分类网络已经足以取得令人满意的分类结果。

SRModule

这里的SRModule是由多个不同复杂度的超分网络分支构成 ${f_{SR}^j \}_{j=1}^M$ 。从广义上来讲，每个分支可以任意的可学习超分网络。由于我们的目标在于：对现有超分方法(如FSRCNN、CARN等)进行加速。因此，我们采用这类轻量型网络作为基线，并设置为最复杂的分支 $f_{SR}^M$ ，其他分支则通过调整通道数来控制网络复杂度。

那么每个分支的通道数该如何确定呢？一个基本原则：对应分支应当取得基线网络在对应类别的性能，可参考下表中的通道与性能的对比。当然，我们也可以通过调整网络的深度来控制模型复杂度，但只要满足上述原则即可。

Classification Method

在训练过程中，ClassModule将对输入子图按照复原难度(而非预定义标签)进行进行分类。因此，与测试阶段不同，在训练阶段每个子图均应经由M个SR分支进行处理。除此之外，为确保ClassModule可以从重建结果接收到回传梯度，我们对每个重建子图以及对应的分类概率相乘得到最终的SR结果：
$\sum_{i=1}^M P_i(x) \times f_{SR}^i(x)$
我们通过调整ImageLoss(即 $L_1$ 损失)来约束重建结果y，我们可以自动得到分类概率。但是在测试阶段，输入子图仅仅经过最大概率的SR分支。因此，我们提出了 $L_c$ 损失(Class-Loss)以使得最大概率接近1，进而确保y以概率1等价于超分子图。

需要注意的是，如果我们仅仅采用ImageLoss与Class-Loss，训练过程很容易收敛到极值点(即最复杂分支)。为避免该现象，我们设计了 $L_a$ （Average-Loss）以约束分类结果。此即为本文所提出的一种新的分类方法。

Loss Functions

损失函数包含三个损失，即Image-Loss、Class-Loss以及Average-Loss。具体来讲，Image-Loss用于确保图像重建质量；Class-Loss用于提升分类的有效性；Average-Loss用于确保每个SR分支可以被均等的选择。总体损失函数定义如下：
$w_1 \times L_1 + w_2 \times L_c + w_3 \times L_a$
在上述三个损失中，Image-Loss很容易理解，即( $L_1$ )；那么另外两个损失如何定义呢？这就是接下来的要介绍的了。先来看一下Class-Loss,我们希望最大概率的类具有更高的置信度，比如 $[0.9, 0.05, 0.05]$ 就比 $[0.34, 0.33, 0.33]$ 更好。其定义如下:
$L_c = -\sum_{i=1}^{M-1} \sum_{j=i+1}^{M} |P_i(x) - P_j(x)|, s.t. \sum_{i=1}^M P_i(x) = 1$
该损失可以极大的提升不同分类结果的概率差异，进而使得最大概率值接近1。正如前面所提到的“仅仅采用Image-Loss与Class-Loss会使得子图倾向于选择最复杂分支”，这是因为复杂分支更容易取得更好的结果，这就是使得ClassModule失其本意，进而SRModule将退化为基线模型。为避免该现象发生，我们应当确保每个SR分支被均等的选择。因此，我们提出了Average-Loss以约束分类结果，定义如下：
$L_a = \sum_{i=1}^M |\sum_{j=1}^B P_i(x_j) - \frac{B}{M}|$
注：B表示batch size。由于统计数值不会进行梯度传播，故而我们采用概率和计算子图数量。通过该损失，训练阶段经过每个SR分支的子图数量近似相同。

Training Strategy

如果我们同时从头开始训练ClassModule与SRModule，性能会非常不稳定，分类模块很容易陷入不好的局部最小。因此，在训练阶段，我们采用三步法：

First：预训练SRModule。为预训练SRModule，我们采用PSNR值进行数据分类。具体来讲，以预训练MSRResNet的PSNR进行划分，前三分之一作为hard类，后三分之一作为simple类。然后，基于所划分的数据集进行不同复杂度的SR分支模型的训练。
Second：固定SRModule，训练ClassModule；
Third：联调ClassModule与SRModule。

Experiments

在数据集方面，我们选择DIV2K，并按照四个尺度(0.6,0.7,0.8,0.9)下采样生成HR图像，这些图像进一步下采样四倍得到LR图像，最后我们裁剪得到1.59M个大小为 $32\times 32$ 的LR子图，这些子图将被均匀的划分为三类。在训练阶段，将通过镜像与旋转进一步增广。除此之外，我们从DIV2K验证集中挑选10张图像(0801-0810)用于训练过程中的验证。

在测试集方面，我们没有选择Set5、Set14这些(分辨率太小)，我从DIV8K中选择300张图像(1201-1500)。具体来讲，前两百分别下采样到2K和4K，称之为Test2K、Test4K，最后的100则称之为Test8K。在测试阶段，LR图像裁剪为 $32\times 32$ 大小，步长为28。

下表给出了FSRCNN、CARN、RCAN三种网络及对应ClassSR的性能对比。很明显，相比基线方法，所提方法可以取得更加性能同时具有更低的计算复杂度，从70%~50%不等。计算复杂度的降低与测试数据高度相关，Test8K可以取得近50%的计算复杂度降低。这是因为更大的图形可以拆分为更多的子图，包含更多的simple子图。