BatchNorm 在 U-Net 中的实战作用 —— 踩过的坑
BatchNorm 简介
BatchNormalization 作为神经网络训练的大杀器,广泛的应用于各类网络的训练中。其好处在于:
- 收敛过程大大加快,从而提升训练速度;
- 提升网络性能,这点可以理解成类似于 Dropout 的防止过拟合的正则化方式;
- 简化调参过程,对参数初始化要求没有那么高,可以使用大的学习率;
对于 BatchNormalization 的理论理解可以阅读 郭耀华’s Blog。其核心问题在于,对于每个隐层神经元,把逐渐向非线性激活函数取值极限饱和区靠拢的输入分布强行拉回到均值为0方差为1的正态分布 (加上 scale 和 shift 操作),使得非线性变换函数的输入值落入对输入比较敏感的区域,以此避免梯度消失问题。
BatchNorm 代码实现 (pytorch)
import torch.nn as nn
self.conv = nn.Sequential(nn.Conv2d(..), # convolution operation
nn.ReLU(),
nn.BatchNorm2d(channels), # channel_num is the Conv2d output channel_num
)
具体操作时,可以使用上述模块将 BatchNorm 和 Convolution 操作集成在一起。有一点值得注意,按照 batch normalization 的原文 1 的解释,需要将其放在 Convolution 后面 激活函数前面;然而实际操作 2 也表明放在 ReLU 后面效果更好。关于 BatchNormalization 位置的讨论可以参考 https://www.zhihu.com/question/283715823/answer/438882036。
实际在 U-Net 中有无 BatchNorm 的影响
最近简单实验了一下 BatchNorm 在经典分割网络 U-Net 中的效果,做了消融性实验:是否在每一次 Conv 之后采用 BatchNorm。实验结果如下:
上图是训练过程中 validation dataset 上的结果,红色曲线代表有 BN 操作的 U-Net,而绿色曲线代表无 BN 操作的 U-Net。尤其是绿色曲线在 2000 次迭代以内验证集上的 dice score = 0,这样的结果很容易让人误认为是算法或者代码本身有问题,而不会考虑到这个现象仅仅是由 BN 操作引发的。 上述实验结果也基本符合理论上 BN 的优势。