0%

Notes on Batch Normalization

在训练深层神经网络的过程中,由于输入层的参数在不停的变化,因此,导致了当前层的分布在不停的变化,这就导致了在训练的过程中,要求 learning rate 要设置的非常小,另外,对参数的初始化的要求也很高。作者把这种现象称为 internal convariate shift. Batch Normalization 的提出就是为了解决这个问题的。BN 在每一个 training mini-batch 中对每一个 feature 进行 normalize. 通过这种方法,使得网络可以使用较大的 learning rate, 而且,BN 具有一定的 regularization 作用。

为什么需要 Batch Normalization

在神经网络的优化中最常用最进本的方法是 SGD, 其目标是寻找最小化 loss function 的参数:

\[ \theta = \underset{\theta}{\mathrm{arg min}} \dfrac{1}{N} \sum_{i=1}^N{\mathcal L\left(x_i, \theta\right)} \]

在求解的过程中,一般是使用 minibatch 的方法,简单来说,就是计算下面的梯度:

\[ \dfrac{1}{m}\sum \dfrac{\partial{\mathcal L \left(x_i, \theta\right)}}{\partial \theta} \]

使用 minibatch 的方法有两个好处:

  1. minibatch 计算出来的 loss 可以看做是整个 trainset 的 loss 的近似值。
  2. minibatch 中,可以并行地计算 m 各样本,因此,使用 minibatch 的方法比原始的 SGD 方法速度更快。

然而,SGD 算法有其固有缺点:

  1. 对初始值要求很高,如果参数的初始化不好,经常不能收敛
  2. 学习率比较难设置,由于每一层 input 数据的 scale 不同,导致 backward 的梯度的 scale 也不同,为了保证不会 gradient vanish, 只能设置较小的 learning rate, 而,较小的 learning rate 使得整个学习过程很慢
  3. 第 N 层的输入受前面 N-1 层的影响,在深度学习中,网络层数很多,因此,及时前面 layer 的很小的影响,当到达第 N 层的时候,会被放大很多倍。

在深度神经网络中,每一层输入数据的分布都不同,因此,每一层的参数都要去学习不同的分布。而主要由于上述 #3 的原因,使得这个过程比较困难。为了说明这个问题,使用一个简单的例子。考虑如下的一个两层的神经网络:

\[ F_2\left(F_1\left(\theta_1, x\right), \theta_2\right) \]

\(F_1\) 的输出作为 \(F_2\) 的输入,因此,在学习的过程中,当 \(F_1\) 的输入 \(x\) 的分布变化的时候 (在选择 minibatch 的使用根本无法保证每次选择的 minibatch 的分布是相同的), \(F_2\) 的参数在向最优解收敛的过程中就会产生偏差,因此,导致了收敛速度变慢。而 BN 就是要解决这个问题。
想象一下,如果 \(F_2\) 的学习不受 \(F_1\) 的输入 \(x\) 的影响,即不管 \(x\) 输入的是什么,\(F_2\) 的输入都是相同的分布,这样,\(F_2\) 就不用调整去适应由于 \(x\) 输入变化带来的影响,那么,是否就可以解决这个问题了呢?因此,这就是 BN 的提出。

Batch Normalization 是什么

Input: Values of \(x\) over a mini-batch, \(\mathcal{B}={x_{1\dots m}}\)
Output: \(\gamma, \beta\)

\[ \begin{align} {\mu}_{\beta} &=\frac{1}{m}\sum_{i=1}^{m}x_i\\ {\sigma}_{\beta}^{2} &=\frac{1}{m}\sum_{i=1}^{m}\left(x_i-{\mu}_{\beta}\right)^2\\ \hat{x}_i &=\frac{x_i-{\mu}_{\beta}}{\sqrt{{\sigma}_{\beta}^2+\epsilon}}\\ y_i &= {\gamma}\hat{x}_i+\beta \end{align} \]

为什么 Batch Normalization 可以加速训练

  1. 允许网络使用较高的 learning rate. 在传统的深度网络训练中,如果使用较大的 learning rate 很容易导致 gradient vanish 或者 gradient explode. 通过在整个网络中 normalize activations, 可以防止参数的较小的改变被应用到较大的或者次优的 activation 中。另一方面,BN 使得网络对于 parameter 的 scale 更加鲁棒。通常情况下,large learning rate 会 increase the scale of layer parameters, 进而会放大 BP 的梯度,导致了 model explosion. BN 的使用使得网络在 BP 的时候不会受到 parameter scale 的影响。这是因为:
$$ \begin{align} & BN\left(Wu\right)=BN\left(\left(aW\right)u\right)\\ \\ & \frac{\partial{BN((aW)u)}}{\partial{u}}=\frac{\partial{Wu}}{\partial{u}} \\ \\ & \frac{\partial{BN((aW)u)}}{\partial{aW}}=\frac{1}{a} \cdot \frac{\partial{Wu}}{\partial{W}} \end{align} $$
  1. 具有一定的 regularization 作用,可以减少 Dropout 的使用。dropout 的作用是方法 overfitting, 实验发现,BN 可以 reduce overfitting.
  2. 降低 \(L_2\) 权重衰减系数。
  3. 取消 LRN(Local Response Normalization).
  4. Reduce the photomatric distortions. 因为 BN 使得训练过程更快,能 observe 到的 sample 次数变少,所以,减少 distorting 使得网络 focus 在真实的图片上面。
  5. BN 不仅仅限定在 ReLU 上,而且,对其它的 activation 也同样适用。

MXNet 代码

Forward 阶段

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
virtual void Forward(const OpContext &ctx,
const std::vector<TBlob> &in_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &out_data,
const std::vector<TBlob> &aux_states) {
using namespace mshadow;
using namespace mshadow::expr;
CHECK_EQ(in_data.size(), 3);
CHECK_EQ(aux_states.size(), 2);
if (ctx.is_train) {
CHECK_EQ(out_data.size(), 3);
CHECK_EQ(req.size(), 3);
​ } else {
CHECK_GE(out_data.size(), 1);
CHECK_GE(req.size(), 1);
CHECK_EQ(req[batchnorm::kOut], kWriteTo);
​ }

Stream<xpu> *s = ctx.get_stream<xpu>();
// 1/m
const real_t scale = static_cast<real_t>(in_data[batchnorm::kData].shape_[1]) /
static_cast<real_t>(in_data[batchnorm::kData].shape_.Size());
Tensor<xpu, 4> data;
Tensor<xpu, 4> out;
if (in_data[batchnorm::kData].ndim() == 2) {
Shape<4> dshape = Shape4(in_data[batchnorm::kData].shape_[0],
in_data[batchnorm::kData].shape_[1], 1, 1);
// 输入数据
data = in_data[batchnorm::kData].get_with_shape<xpu, 4, real_t>(dshape, s);
// 输出数据
out = out_data[batchnorm::kOut].get_with_shape<xpu, 4, real_t>(dshape, s);
} else {
data = in_data[batchnorm::kData].get<xpu, 4, real_t>(s);
out = out_data[batchnorm::kOut].get<xpu, 4, real_t>(s);
}
// slope 是原文中的 \gamma
Tensor<xpu, 1> slope = in_data[batchnorm::kGamma].get<xpu, 1, real_t>(s);
// bias 是算法中的 \beta
Tensor<xpu, 1> bias = in_data[batchnorm::kBeta].get<xpu, 1, real_t>(s);
// 获取 moving average, 如果 use_global_stats, 那么就要使用 moving average
Tensor<xpu, 1> moving_mean = aux_states[batchnorm::kMovingMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> moving_var = aux_states[batchnorm::kMovingVar].get<xpu, 1, real_t>(s);
// whether use global statistics
if (ctx.is_train && !param_.use_global_stats) {
// 不使用 moving average, 那么,要针对当前 minibatch 计算期望和方差
Tensor<xpu, 1> mean = out_data[batchnorm::kMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> var = out_data[batchnorm::kVar].get<xpu, 1, real_t>(s);
CHECK(req[batchnorm::kMean] == kNullOp || req[batchnorm::kMean] == kWriteTo);
CHECK(req[batchnorm::kVar] == kNullOp || req[batchnorm::kVar] == kWriteTo);
// The first three steps must be enforced.
// 计算当前 minibatch 的均值
mean = scale * sumall_except_dim<1>(data);
// 计算当前 minibatch 的方差
var = scale * sumall_except_dim<1>(F<mshadow_op::square>(
data - broadcast<1>(mean, data.shape_)));
if (param_.fix_gamma) {
// 减均值,除以方差,其中 param_.eps 是一个非常小的值,防止出现除 0 的错误
Assign(out, req[batchnorm::kOut], (data - broadcast<1>(mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(var + param_.eps, data.shape_)) +
broadcast<1>(bias, out.shape_));
} else {
Assign(out, req[batchnorm::kOut], broadcast<1>(slope, out.shape_) *
(data - broadcast<1>(mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(var + param_.eps, data.shape_)) +
broadcast<1>(bias, out.shape_));
}
} else {
// test/predict 或者 use_global_stats, 直接使用 moving average
Assign(out, req[batchnorm::kOut], broadcast<1>(slope /
F<mshadow_op::square_root>(moving_var + param_.eps),
data.shape_) * data +
broadcast<1>(bias - (slope * moving_mean) /
F<mshadow_op::square_root>(moving_var + param_.eps), data.shape_));
}
}

Backward 阶段

先给出梯度公式

\[ \begin{align} \frac{\partial{l}}{\partial{\hat{x}_i}} &= \frac{\partial{l}}{\partial{y_i}} \cdot \gamma \\ \frac{\partial{l}}{\partial{\sigma_{\mathcal{B}}^2}} &= \sum_{i=1}^m{\frac{\partial{l}}{\partial{\hat{x}_i}}\cdot\left(x_i-\mu_{\mathcal{B}}\right)} \cdot \frac{-1}{2}\left(\sigma_{\mathcal{B}^2}+\epsilon\right)^{-3/2} \\ \frac{\partial{l}}{\partial{\mu_{\mathcal{B}}}} &= \left( \sum_{i=1}^m{\frac{\partial{l}}{\partial{\hat{x}_i}} \cdot \frac{-1}{\sqrt{\sigma_{\mathcal{B}}^2+\epsilon}}} \right) + \frac{\partial{l}}{\partial{\sigma_{\mathcal{B}}^2}} \cdot \frac{\sum_{i=1}^m{-2(x_i-\mu_{\mathcal{B}})}}{m} \\ \frac{\partial{l}}{\partial{x_i}} &= \frac{\partial{l}}{\partial{\hat{x}_i}} \cdot \frac{1}{\sqrt{\sigma_{\mathcal{B}}^2+\epsilon}} + \frac{\partial{l}}{\partial{\sigma_{\mathcal{B}}^2}} \cdot \frac{2(x_i-\mu_{\mathcal{B}})}{m} + \frac{\partial{l}}{\partial{\mu_{\mathcal{B}}}} \cdot \frac{1}{m} \\ \frac{\partial{l}}{\partial{\gamma}} &= \sum_{i=1}^m{\frac{\partial{l}}{\partial{y_i}}} \cdot \hat{x}_i \\ \frac{\partial{l}}{\partial{\beta}} &= \sum_{i=1}^m{\frac{\partial{l}}{\partial{y_i}}} \end{align} \]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
virtual void Backward(const OpContext &ctx,
const std::vector<TBlob> &out_grad,
const std::vector<TBlob> &in_data,
const std::vector<TBlob> &out_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &in_grad,
const std::vector<TBlob> &aux_states) {
using namespace mshadow;
using namespace mshadow::expr;
CHECK_EQ(out_grad.size(), 1);
CHECK_EQ(in_data.size(), 3);
CHECK_EQ(out_data.size(), 3);
CHECK_EQ(in_grad.size(), 3);
Stream<xpu> *s = ctx.get_stream<xpu>();
Tensor<xpu, 4> data, grad, grad_in;
const real_t scale = static_cast<real_t>(out_grad[batchnorm::kOut].shape_[1]) /
static_cast<real_t>(out_grad[batchnorm::kOut].shape_.Size());
if (in_data[batchnorm::kData].ndim() == 2) {
Shape<4> dshape = Shape4(out_grad[batchnorm::kOut].shape_[0],
out_grad[batchnorm::kOut].shape_[1], 1, 1);
data = in_data[batchnorm::kData].get_with_shape<xpu, 4, real_t>(dshape, s);
grad = out_grad[batchnorm::kOut].get_with_shape<xpu, 4, real_t>(dshape, s);
grad_in = in_grad[batchnorm::kData].get_with_shape<xpu, 4, real_t>(dshape, s);
} else {
data = in_data[batchnorm::kData].get<xpu, 4, real_t>(s);
grad = out_grad[batchnorm::kOut].get<xpu, 4, real_t>(s);
grad_in = in_grad[batchnorm::kData].get<xpu, 4, real_t>(s);
}

Tensor<xpu, 1> mean = out_data[batchnorm::kMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> var = out_data[batchnorm::kVar].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> slope = in_data[batchnorm::kGamma].get<xpu, 1, real_t>(s);
// Tensor<xpu, 1> bias = in_data[kBeta].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> gslope = in_grad[batchnorm::kGamma].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> gbias = in_grad[batchnorm::kBeta].get<xpu, 1, real_t>(s);
// update moving avg
Tensor<xpu, 1> moving_mean = aux_states[batchnorm::kMovingMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> moving_var = aux_states[batchnorm::kMovingVar].get<xpu, 1, real_t>(s);

if (ctx.is_train && !param_.use_global_stats) {
// get requested temp space
Tensor<xpu, 2> workspace = ctx.requested[batchnorm::kTempSpace].get_space<xpu>(
mshadow::Shape2(3, mean.shape_[0]), s);
Tensor<xpu, 1> gmean = workspace[0];
Tensor<xpu, 1> gvar = workspace[1];
Tensor<xpu, 1> tmp = workspace[2];
// 使用 moving average 算法,更新 mean 和 var, 用于test/predict
moving_mean = moving_mean * param_.momentum + mean * (1 - param_.momentum);
moving_var = moving_var * param_.momentum + var * (1 - param_.momentum);
// cal
// 计算方差的梯度
gvar = sumall_except_dim<1>((grad * broadcast<1>(slope, data.shape_)) *
(data - broadcast<1>(mean, data.shape_)) *
-0.5f *
F<mshadow_op::power>(broadcast<1>(var + param_.eps, data.shape_),
-1.5f));
// 计算均值的梯度的第一部分
gmean = sumall_except_dim<1>(grad * broadcast<1>(slope, data.shape_));
gmean *= -1.0f / F<mshadow_op::square_root>(var + param_.eps);
tmp = scale * sumall_except_dim<1>(-2.0f * (data - broadcast<1>(mean, data.shape_)));
tmp *= gvar;
// 均值梯度的第二部分
gmean += tmp;
// assign
if (!param_.fix_gamma) {
//\gamma 的梯度
Assign(gslope, req[batchnorm::kGamma],
sumall_except_dim<1>(
grad * (data - broadcast<1>(mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(var + param_.eps, data.shape_))));
// 输入 x 的梯度
Assign(grad_in, req[batchnorm::kData],
(grad * broadcast<1>(slope, data.shape_)) *
broadcast<1>(1.0f / F<mshadow_op::square_root>(var + param_.eps), data.shape_) +
broadcast<1>(gvar, data.shape_) * scale * 2.0f * (data - broadcast<1>(mean,
data.shape_)) +
broadcast<1>(gmean, data.shape_) * scale);
} else {
Assign(grad_in, req[batchnorm::kData], grad *
broadcast<1>(1.0f / F<mshadow_op::square_root>(var + param_.eps), data.shape_) +
broadcast<1>(gvar, data.shape_) * scale * 2.0f * (data - broadcast<1>(mean,
data.shape_)) +
broadcast<1>(gmean, data.shape_) * scale);
}
// 偏置的梯度
Assign(gbias, req[batchnorm::kBeta], sumall_except_dim<1>(grad));
} else {
// use global statistics with freeze moving mean and var.
if (!param_.fix_gamma) {
Assign(gslope, req[batchnorm::kGamma],
sumall_except_dim<1>(
grad * (data - broadcast<1>(moving_mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(moving_var + param_.eps, data.shape_))));
Assign(grad_in, req[batchnorm::kData], (grad * broadcast<1>(slope, data.shape_)) *
broadcast<1>(
1.0f / F<mshadow_op::square_root>(moving_var + param_.eps), data.shape_));
} else {
Assign(grad_in, req[batchnorm::kData], grad *
broadcast<1>(
1.0f / F<mshadow_op::square_root>(moving_var + param_.eps), data.shape_));
}
}
}

Batch Normalization 使用方法

  1. 一般 BN layer 放在 FC 和 conv 的后面
  2. 在 train 阶段,对每一个 minibatch 使用 BN, 那么,在 test 或者 predict 的时候怎么使用 BN 呢?常见的做法是使用整个 train set 计算出 \(\mu\). 由于 train set 的数据量非常大,计算 \(\mu\) 计算量非常大,所以,经常采用的技术是使用 moving average 算法,在 train 阶段计算出 \(\mu\)

CAFFE 中的使用方法

使用 residual net 网络作为示例

1
2
3
4
5
6
7
8
9
10
11
12
layer {
bottom: "res2a_branch2b"
top: "res2a_branch2b"
name: "bn2a_branch2b"
type: "BatchNorm"
# use_global_stats 表示是否使用全部数据的统计值(该数据实在train 阶段通过
# moving average 方法计算得到)训练阶段设置为 fasle, 表示通过当前的
# minibatch 数据计算得到, test/predict 阶段使用 通过全部数据计算得到的统计值
batch_norm_param {
use_global_stats: true
}
}

MXNet 中的使用方法

1
2
3
4
# data 是输出数据,eps 是为了防止出现除 0 二废除 0 而给出的一个很小的值,momentum
# 是 moving average 的动量项,fix_gamma 参考上面的 C++ 代码,use_global_stats 作用同
# CAFFE 中的解释。
bn = mxnet.symbol.BatchNorm(data, eps, momentum, fix_gamma, use_global_stats)

Reference

Batch Normalization