计图开源：新型归一化方法RBN入选CVPR’21 Oral，性能较PyTorch大幅提升

分享
1

程明明教授带领的南开大学媒体计算研究组的高尚华，韩琦等同学，和香港科大、腾讯的学者合作，提出一种新的新型归一化方法，Representative BatchNorm（RBN)，被CVPR 2021录用为Oral paper，并已在计图开源。

什么是BatchNorm

卷积神经网络（CNNs）以其强大的表征能力提高了各种计算机视觉任务的性能。随着结构复杂性的增加和模型参数量的增长，CNN的训练变得愈发困难。BatchNorm[1]通过在训练过程中不断归一化中间层特征，来使得网络的训练变得更加稳定，因而被广泛应用于各种神经网络架构中。

BatchNorm使用batch维度的统计均值和方差归一化CNN的中间层特征，但同时不可避免地会忽略不同样本之间的表征差异。由于训练过程中的batch维度的统计信息与测试样本的统计信息之间存在一定差异，会导致模型对测试样本进行特征提取时引入额外的噪声或忽略重要特征，从而无法实现最优的表征效果。

部分工作利用基于实例的统计信息来摆脱batch维度依赖造成的副作用。例如YuxinWu 和Kaiming He等人提出的GroupNorm[2]使用单个样本的多通道特征统计值实现归一化。

但由于缺少batch信息，训练的不稳定导致其性能在多数情况下难以与BatchNorm媲美。另一部分工作通过结合多种维度的统计信息进行归一化，但往往会引入更多的计算开销。

Representative BatchNorm

程明明教授及其学生高尚华等在CVPR2021 Oral 论文中，提出了一种在保留BatchNorm优势的同时，以极低计算开销对样本特征进行矫正的新型归一化方法，Representative BatchNorm（RBN）[3]。该方法是在BatchNorm的中心化和缩放操作中分别添加了简单而有效的中心化校准和缩放校准。RBN 公式如下：

中K_m 和 K_s分别为样本统计信息，可根据任务定义；w_m，w_v 和 w_b为可学习参数，仅需要在每个通道中引入三个权重，即可自适应学习特征的矫正规则。

中心化校准可增强有效特征（图1a）并减少噪音特征（图1b）。缩放校准限制了特征强度以形成更稳定的特征分布（减小图1c，1d两个通道的特征强度差异）。此外，作者也从理论上证明了这两种校准的有效性，详见论文[3]。

图1 RBN的中心化校准和缩放校准可视化对比

为什么Jittor会大幅提升RBN的性能呢？

虽然RBN引入的计算量可忽略不计。但是由于现有深度学习平台并没有对其进行优化，RBN实现起来还是会引入一定的运算额外时延。得益于 Jittor框架“元算子融合”和“统一计算图”两项新技术，RBN在无需任何优化的情况下运行速度相比于PyTorch框架就有大幅提升。

图2 ResNet101 + RBN在pytorch 与jittor平台的推理速度对比

RBN在各类视觉任务中的应用

RBN可以替换现有方法下的BatchNorm操作，以可忽略的额外计算成本和参数来提高各种任务（例如分类，检测和分割）的性能。

分类 在分类任务下，RBN相较于之前的归一化工作有明显性能优势。RBN也可以进一步提升例如Res2Net[4] 等SOTA的网络结构的性能。

图3 RBN提升分类网络的性能

目标检测与实例分割 在任务下在MaskRCNN +BatchNorm[5]的基线下，MaskRCNN+RBN有明显的性能优势。

图4 RBN应用于检测和示例分割的效果

全景分割 相较于Panoptic PFN+BatchNorm[6]的基线，RBN在全景分割任务下也有显著的性能提升。

图5 RBN应用于全景分割的效果

更多关于RBN的信息，可以参见项目主页

和代码主页

在Jittor框架的加持下，RBN能够在以上任务中实现更优的效率。

参考文献

  1. Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICMLW, 2015.
  2. Yuxin Wu and Kaiming He. Group normalization. In ECCV, 2018, 1-16.
  3. Shang-Hua Gao, Qi Han, Duo Li, Pai Peng, Ming-Ming Cheng and Pai Peng, Representative Batch Normalization with Feature Calibration. In CVPR, 2021.
  4. Shang-Hua Gao, Ming-Ming Cheng, Kai Zhao, Xin-Yu Zhang, Ming-Hsuan Yang, Philip Torr, Res2Net: A New Multi-scale Backbone Architecture. IEEE TPAMI, 2001, Vol. 43, No. 2, 652-662.
  5. Kaiming He, Georgia Gkioxari, Piotr Dollar, and Ross Girshick. Mask R-CNN. In ICCV, 2017, 2961-2969.
  6. Alexander Kirillov, Ross Girshick, Kaiming He, and Piotr Dollar. Panoptic feature pyramid networks. In CVPR, 2019, 6399-6408.