RepSR让BN重回SR怀抱,加速训练效率,同时消除伪影问题

2022-06-24 14:16:07 浏览数 (1)

本文对重参数设计理念进行了重审视,对重参数超分网络的关键成分进行了调查。我们发现:BN层有助于引入训练非线性能力并改善模型性能(这与OREPA一文的出发点相同)。但是,由于BN会恶化超分性能并引入伪影问题,所以超分模型往往会忽视BN层。

本文对BN存在的问题进行了仔细分析并提出一种简单有效的解决方案。具体来说,我们首先像往常一样训练带BN的超分网络,然后将BN切换到frozen模式进行后续训练。除了再次将BN引入超分网络外,我们进一步设计了一种新的重参数模块并将所得模型称之为RepSR。

实验结果表明:相比其他重参数超分方案,所提RepSR可以取得更优性能。此外,RepSR取得了更佳的性能-推理耗时均衡

1出发点

本文的出发点主要有这样两个:

  • 尽管Plain架构可以更好的利用GPU并行能力与硬件加速,但是Plain架构的性能通常弱于多分支架构;
  • 重参数技术可以在训练结构利用多分支结构达成更优性能,在推理阶段利用Plain结构的高效率能力。但是,重参数技术的核心BN对于SR任务并不友好:不仅会噪声性能下降,同时会引入伪影问题。

2本文方案

在该部分内容中,我们主要对RepSR进行描述,包含BaseNet(见上图a)以及所提重参数模块(见上图b)。为保持低计算复杂度与内存占用,BaseNet是一种VGG风格网络:主体部分由多个3times 3 卷积堆叠而成(激活函数为PreLU),在尾部采用PixelShuffle进行上采样。

如果直接对上述BaseNet进行训练,那么所得方案性能会弱于多分支架构(如EDSR)。重参数技术是有效提升模型性能的trick,为消除剔除BN带来的性能损失,ECBSR设计了一种包含Sobel与Laplacian的重参数模块

不同于ECBSR,本文重新将BN引入到超分中并构建了RepSR模块:

  • Clean Residual Path:不同于RepVGG,该分支并未引入BN层;
  • Expand-and-Sequeeze Convolution:它采用3times 3 卷积扩展通道维度,然后采用1times 1 卷积收缩通道维度,即Conv(3times 3)-BN-Conv(1times 1)
  • Frozen BN Layer:该技术在high-level任务应用较多(如预训练骨干在检测模型中的迁移),但鲜少用于超分任务。相比high-level中的方案,这里主要有这么几点区别:(1) high-level任务在大多数训练过程中均采用FrozenBN,而本文仅在训练的后期(<10%)采用FrozenBN; (2) high-level中使用FrozenBN是为了将预训练模型在下游任务迁移或者因为小batch训练,而本文则是为了利用其训练非线性能力。

如前文所提到,ECBSR同样将重参数技术引入到超分中以达成高效推理。相比ECBSR,RepSR有这样三点优势:

  • RepSR的重参数模块表明:ECBSR中复杂的Sobel与Lapacian并非必要的。RepSR将BN再次引入SR,在享受其训练非线性能力同时未受其弊端影响。
  • 由于复杂的结构与优化问题,ECBSR的性能伴随网路架构加深出现了性能下降(如>16);而RepSR在tiny与large模型方面均具有非常好的性能。
  • 相比ECBSR,RepSR具有更快的训练速度,可节省50%训练耗时。

3重新审视超分中的BN问题

在该部分内容中我们精心分析了BN在超分任务中所引起的伪影问题,然后提出了一种有效策略以解决该问题。下式为BN的公式,在训练阶段均值与方差通过batch数据计算得到,而在推理阶段则采用训练集上估计的统计信息。

y = frac{x - mu}{sqrt{sigma^2 epsilon}} times gamma beta

如下图所示,BN对超分任务不够友好:会产生伪影问题,进而导致PSNR指标下降。为补偿移除BN带来的优化问题,residual scaling、intermediate supervision以及longer training成为了常用的训练策略。

我们发现:BN导致的伪影往往出现在平滑区域。这些区域明显具有不同的局部方差,这就意味着:伪影与BN中的不正确统计信息相关。有意思的是,如果我们采用输入图像的统计信息进行归一化,该伪影问题就可以极大程度的消除。这说明:BN导致的伪影与train-test不一致有关。此外,我们发现:BN伪影问题发生在"Patch"级不一致(可参考下图示例)。这种"Patch"级不一致还会引入另一个问题:验证集无法判断当前训练网络的"健康"状态

需要指出的是:BN伪影仅发生在推理阶段,而训练过程健康且稳定。那么,消除train-test不一致的一个很直观的方案就是:在训练阶段使用population统计信息。因此,我们首先采用mini-batch统计信息训练超分网络,然后再训练后期采用population统计信息进行训练。 通过使用该训练策略,我们不仅可以受益于BN的快速收敛能力,同时可以避免BN伪影问题

需要注意的是,不同于high-level中采用的FrozenBN,我们并未对BN的affine参数进行冻结,而仅对统计参数进行冻结

4Experiments

上表给出了TinySR方面的性能对比,从中可以看到:

  • 重参数方案能够以推理时更低的参数量、更低计算复杂度达成更优的性能;
  • 相比ECBSR,所提RepSR取得类型相当或更优的性能
  • 相比RepSR,由于未使用BN,SESR方案出现了显著性能下降。

上表从训练耗时角度对比了ECBSR与RepSR,可以看到:RepSR具有更快的训练速度

上表&图对比了不同重参数方案的性能,可以看到:

  • 在更大容量模型方面,ECBSR的表现出现了变差;
  • RepSR可以取得比RepVGG与ECBSR更优的性能,甚至取得比EDSR更优的性能;
  • 相比其他方案,RepSR可以重建更清晰的线条。

0 人点赞