Happy · 2022年05月13日

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

3113edd55ed9d734bcbae8969d5261c5.jpg

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

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

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

f0b605ff9940065c704944af311c31ac.jpg

1出发点

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

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

2本文方案

345a93b1c540115fd1a583788417ce89.jpg

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

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

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

  • Clean Residual Path:不同于RepVGG,该分支并未引入BN层;
  • Expand-and-Sequeeze Convolution:它采用卷积扩展通道维度,然后采用卷积收缩通道维度,即;
  • 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数据计算得到,而在推理阶段则采用训练集上估计的统计信息。

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

9efdfe134d3f96f56c6f16b3ce3c1b52.jpg

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

75c5641dfeb34424a1896af7ffb5b7e2.jpg

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

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

4Experiments

91880eecd9207d1ce7eb02fb16d3261e.jpg

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

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

7fdedac287ab351ab67b2d783cab4fcc.jpg

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

8285f7de670f7a6b3eb05399b0462ca1.jpg

5ffce57e8765d2081fceaaa27ffa10f9.jpg

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

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

推荐阅读

本文章著作权归作者所有,任何形式的转载都请注明出处。更多动态滤波,图像质量,超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
6193
内容数
190
夯实深度学习知识基础, 涵盖动态滤波,超分辨,轻量级框架等
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息