极市导读
本文使用了一种训练策略,即:首先像正常的 BN 一样,用小 Batch 来训练 SR 网络。在训练的末期,使用统计得到的均值和方差。使用这种方式训练,既享受与 BN 训练更快收敛的好处,同时不受 BN 伪像的影响。 >>加入极市CV技术交流群,走在计算机视觉的最前沿
本文目录
14 RepSR:解决超分模型的 BN 不友好问题
14.1 VGG 式模型的优势
14.2 结构重参数化方法简介
14.3 从 RepVGG 到 RepSR
14.4 RepSR 基本模型架构
14.5 RepSR 结构重参数化方法
14.6 重新思考超分任务中的 BN 机制:BN 伪像的形成原因分析
14.7 重新思考超分任务中的 BN 机制:去除 BN 伪像的一种简单策略
14.8 RepSR 实验结果
-
14 RepSR:解决超分模型的 BN 不友好问题
论文名称:RepSR: Training Efficient VGG-style Super-Resolution Networks with Structural Re-Parameterization and Batch Normalization
论文地址:
https://arxiv.org/pdf/2205.05671.pdf
14.1 VGG 式模型的优势
结构重参数化方法最早是有丁霄汉大神的 RepVGG[1]首先提出。
结构重参数化方法一作解读 (RepVGG):
https://zhuanlan.zhihu.com/p/344324470
结构重参数化方法用于剪枝 (ResRep[2]):
https://zhuanlan.zhihu.com/p/402106979
结构重参数化方法用于设计 MLP 网络 (RepMLPNet[3]):
https://zhuanlan.zhihu.com/p/375422742
结构重参数化方法用于设计大核卷积网络 (RepLKNet[4]):
https://zhuanlan.zhihu.com/p/481445076
"VGG 式" 模型指的是:
1. 拥有像 VGG 那样的 plain 或 feed-forward 架构,没有任何分支,意味着每一层都以它上一层的输出为输入,并且输出只传递给下一层。
2. 仅使用 3x3 卷积和 ReLU 作为激活函数。
3. 模型的详细架构 (包括深度和通道数) 不经过架构搜索或者人工微调等过程。
RepVGG 模型的基本架构:将 20 多层 3x3 卷积堆起来,分成 5 个 stage,每个 stage 的第一层是 stride=2 的降采样,每个卷积层用 ReLU 作为激活函数。
3×3 卷积的优势:
3×3 卷积的优势是已经被现代计算库 (如 NVIDIA cuDNN[5]) 进行了高度的优化。下图 1 是 NVIDIA 1080Ti 上不同大小卷积核的理论计算量,运行时间和理论的计算密度 (Batch size=32,Input channel=output channel=2048,resolution=56×56,stride=1,cuDNN 7.5.0)。TFLOPS 是指 Tera FLoating-point Operations Per Second,即某种操作的计算密度。3×3 conv 的理论计算密度约为其他的 4 倍,这表明理论的运算量 FLOPs 并不能代表不同架构的实际速度。
"VGG" 这种极简架构的优势是:
-
运算速度快: VGG-16 是 EfficientNet-B3 计算量的 8.4 倍,但是在 1080 Ti 的运行速度却是后者的 1.8 倍,这意味着前者的计算密度是后者的约 15 倍。运行速度的提升除了受益于 Winograd 算法本身之外,还得益于内存访问成本 (Memory Access Cost, MAC) 和并行度。MAC 表示内存使用量,即运行这个操作占用多少内存。比如多分支操作 Inception,ResNet 的 MAC 很大。比如尽管多分支操作所需的计算量可以忽略,但是 MAC 特别高。此外,高并行度的模型在相同的 FLOPs 时比低并行度的模型速度快。 -
节约内存占用: 多分支架构之所以是内存低效的,是因为在计算多路分支输出的过程中,在没有把多路的计算结果最终加起来之前,各个分支的结果都要存在内存里面,导致了较高的内存占用。下图 2 说明这个问题,加了残差连接之后,内存占用峰值变为了原来的 2 倍 (假设每一层的特征大小不变)。 -
灵活性: 单路架构灵活性更好。残差连接要求最终输出的张量维度与输入是同等大小,这就为模型的设计带来了限制。此外,plain 的架构容易改变各层的宽度 (如剪枝) 。但是多分支模型不容易对每层进行剪枝,也会导致显著的性能下降或低加速比。相比之下,简单的架构允许我们灵活地配置每个 conv 层以获得最佳的性能-效率 trade-off。
14.2 结构重参数化方法简介
结构重参数化方法的思想是借助多分支架构和单分支架构各自的优势。多分支架构 (如 ResNet,Inception,DenseNet,各种 NAS 架构) 和单分支结构 (VGG 式模型) 相比各有什么优劣呢?
-
多分支架构训练更加稳定且容易,比如 ResNet 的残差连接就能够很好地解决 gradient vanishing 的问题。但是推理速度慢,占用的内存大。 -
单分支结构推理速度快,节省内存。但是训练比较困难,训练的性能较低。
结构重参数化具体的做法如下图 3 所示。
在训练的时候如图 3 (b) 所示,为每个 3×3 的卷积层添加平行的 1x1 卷积分支和恒等映射分支,构成一个 RepVGG Block。这种设计是借鉴 ResNet 的做法,区别在于 ResNet 是每隔两层或三层加一分支,而作者是每层都加。加完之后模型就变成了个多分支的模型,作者利用多分支架构训练的稳定性来训练这个多分支架构。
在部署的时候如图 3 (c) 所示,使用重参数化方法拿着图 3 (b) 的权重,对模型作等价变换,直接变成图 3 (c) 的 VGG 式架构。
这个等价转换的策略就如下图 4 所示。利用的原理是卷积操作的齐次性,具体而言:
-
1x1 卷积是相当于一个特殊 (卷积核中有很多 0) 的 3x3 卷积。 -
恒等映射是一个特殊 (以单位矩阵为卷积核) 的 1x1 卷积,因此也是一个特殊的 3x3 卷积。
所以我们只需要:
1. 把 1x1 卷积等价转换为 3x3 卷积,多余的部分使用 0 填充。
2. 把 Identity 操作转换为 1x1 卷积,只要构造出一个以单位矩阵为卷积核的 1x1 卷积。
在图 4 中作者展示了 RepVGG Block 结构重参数化的过程。为了便于可视化,取 ,所以 3×3 卷积和 1×1 卷积分别包含 4 个 3×3 矩阵和 4 个 1×1 矩阵。
下面要把每个分支的 BN 层融入卷积的权重里面 (也就是通常所谓的 "吸BN" 操作)。假设这三个分 支中, 卷积的权重是 卷积的权重是 , Identity 操作没有权重。3种不同操作的BN参数分别是:
其中, 分别代表 BN 层的滑动均值,方差, 学习到的缩放因子和偏置。
则输出特征是:
而 BN 操作可以等价为:
把 BN 的参数合并到卷积中的过程为:
合并之后卷积就成为了:
所以即便是带有 BN 层的卷积, 也可以通过以上的操作把 BN 的参数 融入到卷积层的参数 中。以上操作可以使用下图 4 表示。作者把 Identity 操作分支先转换成 卷积 (其参数是单位矩阵), 再将两个 卷积核用 0 给 padding 成 卷积。最后把三个卷积的卷积核和 bias 分别相加得到最终的 RepVGG Block 的 卷积。
14.3 从 RepVGG 到 RepSR
超分模型的前向推理速度是现实世界超分应用的瓶颈。尽管当前的工作已经提出了许多具有较少 FLOPs 的轻量级网络设计,但它们在实际应用中可能不会运行得更快,因为 FLOPs 不能精确反映实际速度,如上文所论述的那样,VGG-16 是 EfficientNet-B3 计算量的 8.4 倍,但是在 1080 Ti 的运行速度却是后者的 1.8 倍,这意味着前者的计算密度是后者的约 15 倍。而且,超分模型中广泛存在的 skip 和 Dense 连接会引入额外的 GPU 内存开销。plain 类的模型能够充分利用 GPU 的并行计算加速能力。
因此,在本文中作者想要设计 plain 类型的超分模型,但是直接训练这样的模型总会得到比较差的性能,因此作者希望借助重参数化技术,在超分模型训练时使用多分支架构,训练好以后转化成 plain 的模型用于部署和推理。具体来说,RepSR 使用多个分支的线性组合来训练模型,并在推理阶段将它们合并为标准的 3×3 卷积。训练效率提高了,但是没有牺牲推理速度。为了充分利用多分支结束,有必要在每个分支中添加 BN 层。BN 层可以在训练过程中引入非线性,提高最终性能。其他激活函数很难取代 BN,因为它们不能合并成一个卷积。
但是,BN 操作对于超分任务而言并不友好,本文对超分模型的 BN 不友好问题进行了分析,认为 BN 层之所以会影响超分模型的性能以及产生不好看的伪影的原因是:训练-测试不一致性。有趣的是,作者的分析还表明,这个 BN 问题只出现在推理过程中,而训练过程是 "稳定" 且 "健康" 的。为了解决这个问题,本文使用了一种训练策略,即:首先像正常的 BN 一样,用小 Batch 来训练 SR 网络。在训练的末期,使用统计得到的均值和方差。使用这种方式训练,我们既享受与 BN 训练更快收敛的好处,同时不受 BN 伪像的影响。
14.4 RepSR 基本模型架构
RepSR 基本模型架构如下图 5 (a) 所示,为了保持较低的计算成本和内存消耗,RepSR 推理时包含一系列的 3×3 卷积和激活函数 (PReLU)。在模型的末尾使用 PixelShuffle 操作和最近邻插值的上采样操作。
RepSR 基本模型架构和 ECBSR[6] 保持一致。该结构已经被 ECB 证明比其他复杂的网络结构更易于部署。根据在 ECBSR 中的测量,具有 plain 网络结构的 FSRCNN[7] 比具有复杂拓扑结构的 IMDN-RTC[8] 具有更高的 FLOPs,但在 SnapDragon 865 上运行速度快两个数量级,可实现 540p 至 1080p 的扩展。
BN 操作对于超分任务而言并不友好,ECBSR 方法通过手工设计的复杂模块 (包含 Sobel 和 Laplacian 滤波器) 来缓解这一问题。RepSR 块包含一个 clean residual 支路,两个带有 Frozen BN 层的 Expand-and-Squeeze 支路。
Clean Residual 支路: RepVGG 这个工作中 Identity 分支也是使用 BN 层的,但是 RepSR 发现不使用 BN 层的 Identity 分支在超分任务里面可以达到更好的性能。主要是由于 low-level 的 SR 任务和 high-level 任务的信息处理机制不同。SR 模型中后一个块增强了前一个块的特征,而不是像那些 high-level 任务那样将它们转换到更高级的语义空间中。因此作者需要无 BN 的 shortcut 帮助更好地传播信息。
Expand-and-Squeeze Convolution 支路: 作者对其他两个分支采用 Expand-and-Squeeze Convolution。具体而言,首先使用 3×3 卷积来扩展通道数,然后使用 1×1 卷积来压缩通道数,即 3×3 Conv-BN-1×1 Conv。
式中, 是 卷积的权重, 是 1×1 卷积的权重。这样中间宽,两边窄的设计意图是增加模型的表达能力。另外,作者发现使用了 BN 操作之后,模型的训练过程更加稳定。
Frozen BN layer 训练策略: 本文使用了一种训练策略,即:首先像正常的 BN 一样,用小 Batch 来训练 SR 网络。在训练的末期,使用统计得到的均值和方差。这个策略在高级任务中被命名为 FrozenBN,虽然使用了相同的术语,但它们在三个方面有所不同:
1) FrozenBN 常用于 high-level 的任务,本文首次在超分任务中使用。
2) FrozenBN 常用于把预训练模型迁移到下游任务中进行训练,但本文用于 SR 任务的重参数化。
3) 与之前的 FrozenBN 在大部分训练期间 BN 层被 Frozen 不同,本文的 FrozenBN 在大部分训练期间 BN 层被训练,并且仅在训练结束时被 Frozen 非常短的时间 (例如<10%)。
与 ECBSR 的区别: ECBSR 是第一个将重参数化方法用于 SR 任务的技术,RepSR 与 ECBSR 的区别是:
-
RepSR 不采用 ECBSR 中采用的一些复杂的 Sobel 和 Laplacian filters。这主要是因为 RepSR 将 BN 重新引入到 SR 中,使得在训练期间享受来自 BN 的非线性的好处,同时不受 SR 中 BN 伪像的影响。 -
ECBSR 的性能对于更深的网络 (>16 个卷积层) 急剧下降。而 RepSR 在小型和大型网络上都可以获得良好的性能。 -
由于引入了 BN 操作,RepSR 比 ECBSR 具有更快的训练速度。比如可以减少 50%的训练时间。
与 RepVGG 的区别: RepVGG 使用重参数化技术在 high-level 的视觉任务上取得了极佳的性能。但是:
-
直接将 RepVGG 的技术引入 SR 任务会得到较差的性能,因为 BN 对超分任务并不友好,且会造成人工伪影。为了解决 BN 问题,本文提供了 BN 伪影的详细分析,并提出了 FrozenBN 策略。 -
RepSR 提出了 Clean Residual 支路,Expand-and-Squeeze Convolution 支路等用于 SR 任务的特殊设计。
-
14.5 RepSR 结构重参数化方法
第 1 步:吸 BN,把 BN 层参数融合到前置位卷积中。
第 2 步:融合 3×3 卷积和 1×1 卷积:
式中, 是 卷积的权重, 是 1×1 卷积的权重。根据卷积的结合律,能够得到 merged weight 和 merged bias。因为第二个卷积是 1×1 卷积,只进行 channel 维度的变换,所以这一步可以通过 batch matrix multiplication 完成。
整个算法如下图 6 所示。
14.6 重新思考超分任务中的 BN 机制:BN 伪像的形成原因分析
BN 方法的原理可以写成:
其中, 是可学习的仿射参数, 是一个非常小的数值。在训练过程中, 计算每个 channel 的 , 在推理过程中, 使用统计值。
然而,如图 7 所示,BN 对 SR 任务是不友好的。它倾向于在推断过程中引入令人不愉快的伪像。这种伪像表现为输出中混杂着不自然的色彩,导致 PSNR 大幅下降。因此,在诸如 ESRGAN[9] 的一些经典工作中一般都将 BN 移除。
要回答这个问题,作者首先从一个观察开始,即:BN 产生的伪像更可能出现在图像的平滑区域。
我们知道,BN 层的统计信息 对于每个 channel 而言都是相同的,如下式所示:
既然对于图片中相同的 channel 而言,不同位置的 是一样的,那么对于产生伪像的原因,作者提出了假设:
一张图片 的不同 patch 的局部统计信息是不一样的, 但是不同 patch 的 BN 统计信息 确是相同的, 这样的 gap 导致了伪像的产生, 也就说明了 BN 伪像在 patch level 产生。
在一张图片内部各个 patch 的 feature 的统计方差不一致,导致这些位置的超分结果存在伪像。如上图 7 的第 3 列所示,如果将推理时的统计信息 换成是单张输入图片的信息,则伪像会大大减少。这一现象意味着 BN 伪像的原因可能来自训练测试的不一致性,即小批量统计和总体统计之间的不一致性。
BN 伪像在 patch level 产生。为了验证这一猜想,作者进行了一个实验:如下图 8 所示,作者将 LR 图片的某一块 (这块的超分结果会产生 BN 伪像) 给粘贴在另一张图片的中心位置 (中心位置的原始超分结果没有伪像) 上,并把这张合成图片通过超分模型。结果显示,BN 伪像仍然存在于合成图片上,而且伪像的样子和原来非常相似。这意味着 SR 模型非常关注一个块内部自身的统计信息,BN 伪像在 patch level 产生。那些 "有问题的",会产生伪像的块,无论是单独出现,还是和别的块一起,都依旧会产生伪像。
另外作者观察了有伪像的位置,BN 层统计信息,和无伪像的位置的方差的 范数曲线。观察到无伪像的位置的方差量级和 BN 层统计方差量级接近,但是有伪像的位置的方差量级和 BN 层统计方差量级差距拉大。因此在除以方差时由于不正常的方差的量级导致了伪像的形成。
14.7 重新思考超分任务中的 BN 机制:去除 BN 伪像的一种简单策略
如下图 10 所示,作者首先观察到 training loss 是正常,稳定地下降的,但是 validation performance 非常抖动。 作者基于上述的分析,认为解决 BN 伪像问题的一种方法是减小训练-测试的差距。
具体来讲,我们把关注点放在 BN 的统计信息 \mu,\sigma\mu,\sigma 上面:
-
在训练阶段, 是不断变化的。 -
在测试阶段, 使用训练阶段的统计信息。
而训练阶段 和测试阶段的 统计信息之间的 gap 就是造成 BN 伪像的原因。
为此,作者提出了一种策略,也就是 FrozenBN 的改进版,即:在训练的初始阶段, 正常更新,但是在训练的最后阶段 (例如<10%), 被固定住,使用训练阶段的统计信息继续进行训练。通过这种策略来缩小训练-测试的差距。
与 FrozenBN 不同,作者不冻结 BN 中的仿射参数 ,因为这是不必要的。调整仿射参数可以得到更大的优化空间来去除 BN 伪影。
14.8 RepSR 实验结果
训练集使用 DIV2K,验证集使用 Set5,Set14,BSD100,Urban100,DIV2K 验证集,使用 L1 Loss 训练,优化器使用 Adam,1000K iterations,batch size 设为 32,将 RGB 转换为 YCbCr 格式后在 Y 通道上计算 PSNR 和 SSIM。
作者将 RepSR 与具有代表性的 SR 方法进行比较,包括 SRCNN,FSRCNN,ESPCN,VDSR,LapSRN,CARN ,TPSR,IMDN,以及最近的 SR 重参数化工作:ECBSR 和 SESR。MxCy 代表重参数化之后有 x 个卷积,y 个 channel。
如图 11 所示是 RepSR 相关的实验结果,可以得出以下结论:
-
重参数化的方法可以用更少的参数和更低的推理计算代价来提高性能。 -
与 ECBSR 相比,RepSR 具有简单的结构,实现了相当或更好的性能,而 ECBSR 需要更复杂的结构和手动设计的滤波器。 -
SESR 有很大的性能下降,表明 BN 在重新参数化中的重要性。
作者进一步加大了网络的尺寸,将 RepSR 与 SwinIR-Light,RCAN,EDSR,ECBSR,IMDN 和 RFDN 等模型对比,结果如下图 12 所示为 PSNR 和 Throughput 的对比。吞吐量在 A100 上测试,batch 为 4,fp16 精度,上采样范围为 320×180 至 1280×720。使用 Y 通道在 Urban100 测试集上进行评估。RepSR 能够在性能和吞吐量之间取得更好的平衡。
下图是超分的可视化结果,RepSR 可以恢复更清晰的线条,甚至产生比 EDSR 更好的结果。
总结
本文中作者想要借助重参数化技术设计 plain 类型的超分模型。RepSR 使用多个分支的线性组合来训练模型,并在推理阶段将它们合并为标准的 3×3 卷积。训练效率提高了,但是没有牺牲推理速度。为了充分利用多分支结束,有必要在每个分支中添加 BN 层。为了解决 BN 在 SR 任务中的不友好问题,作者对这个问题进行了分析,认为 BN 层之所以会影响超分模型的性能以及产生不好看的伪影的原因是:训练-测试不一致性。有趣的是,作者的分析还表明,这个 BN 问题只出现在推理过程中,而训练过程是 "稳定" 且 "健康" 的。为了解决这个问题,本文使用了一种训练策略,即:首先像正常的 BN 一样,用小 Batch 来训练 SR 网络。在训练的末期,使用统计得到的均值和方差。使用这种方式训练,RepSR 既享受与 BN 训练更快收敛的好处,同时不受 BN 伪像的影响。
参考
-
^RepVGG: Making VGG-style ConvNets Great Again -
^ResRep: Lossless CNN Pruning via Decoupling Remembering and Forgetting -
^RepMLPNet: Hierarchical Vision MLP with Re-parameterized Locality -
^Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs -
^ cudnn: Efficient primitives for deep learning -
^Edge-oriented Convolution Block for Real-time Super Resolution on Mobile Devices -
^Accelerating the super- resolution convolutional neural network -
^ Lightweight image super-resolution with information multi-distillation network -
^ESRGAN: Enhanced super-resolution generative adversarial networks

公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT
极市干货
# 极市平台签约作者#
科技猛兽
知乎:科技猛兽
清华大学自动化系19级硕士
研究领域:AI边缘计算 (Efficient AI with Tiny Resource):专注模型压缩,搜索,量化,加速,加法网络,以及它们与其他任务的结合,更好地服务于端侧设备。
作品精选

