深度學習網絡層之 Batch Normalization

Batch Normalization

S. Ioffe 和 C. Szegedy 在2015年《Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift》論文中提出此方法來減緩網絡參數初始化的難處.

Batch Norm原理

內部協轉移（Internal Covariate Shift）：由於訓練時網絡參數的改變導致的網絡層輸出結果分布的不同。這正是導致網絡訓練困難的原因。

對輸入進行白化（whiten：0均值，單位標准差，並且decorrelate去相關）已被證明能夠加速收斂速度，參考Efficient backprop. （LeCun et al.1998b）和 A convergence anal-ysis of log-linear training.（Wiesler & Ney,2011）。由於白化中去除相關性類似PCA等操作，在特征維度較高時計算復雜度較高，因此提出了兩種簡化方式：
1）對特征的每個維度進行標准化，忽略白化中的去除相關性；
2）在每個mini-batch中計算均值和方差來替代整體訓練集的計算。

batch normalization中即使不對每層的輸入進行去相關，也能加速收斂。通俗的理解是在網絡的每一層的輸入都先做標准化預處理(公式中k為通道channel數),再向整體數據的均值方差方向轉換.
標准化:

\[\text{unit gaussian activations} \\ \hat x^{(k)}=\frac{x^{(k)}-E[x^{(k)}]}{\sqrt{\operatorname{Var[x^{(k)}]}}} \]

Batch Normalization轉換式子:

\[\text{Batch Normalizing Transform} \\ \begin{align} \mu &= \frac{1}{m}\sum_{i=1}^m x_i &\text{//mini-batch mean} \\ \sigma^2 &= \frac{1}{m} \sum_{i=1}^m (x_i - \mu)^2 &\text{//mini-batch variance} \\ \hat{x_i} &= \frac{x_i - u}{\sqrt{\sigma^2 + \epsilon}} &\text{//normalize} \\ y_i &= \gamma \hat{x_i} + \beta &\text{//scale and shift} \end{align} \]

其中ϵ是為了防止方差為0導致數值計算的不穩定而添加的一個小數，如1e−6。

Batch Norm特征轉換scale

隨着前邊各層的累計影響,導致某一層的特征在非線性層的飽和區域,因此如果能對該特征做變換,使之處於較好的非線性區域,那么可使得信號傳播更有效.假如將標准化的特征直接交給非線性激活函數,如sigmoid,則特征被限制在線性區域，這樣改變了原始特征的分布。

論文中引入了兩個可學習的參數\(\gamma , \beta\)來近似還原原始特征分布。這兩個參數學習的目標是\(\gamma=\sqrt{\text{Var}[X]}、 \beta=\mathbb E[X]\),其中\(X\)表示所有樣本在該層的特征.原來的分布方差和均值由前層的各種參數weight耦合控制，而現在僅由\(\gamma , \beta\)控制，這樣在保留BN層足夠的學習能力的同時，使其學習更加容易。因此，加速收斂並非由於計算量減少（反而由於增加了參數增加了計算量）。

那么為什么要先normalize再通過\(\gamma , \beta\)線性變換恢復接近原來的樣子，這不是多此一舉嗎？
在一定條件下可以糾正原始數據的分布（方差，均值變為新值γ,β），當原始數據分布足夠好時就是恆等映射，不改變分布。如果不做BN，方差和均值對前面網絡的參數有復雜的關聯依賴，具有復雜的非線性。在新參數 γH′ + β 中僅由 γ,β 確定，與前邊網絡的參數無關，因此新參數很容易通過梯度下降來學習，能夠學習到較好的分布。

反向傳播

反向傳播梯度計算如下：

BN的前向與反向傳播示意圖:

在訓練時計算mini-batch的均值和標准差並進行反向傳播訓練，而測試時並沒有batch的概念，訓練完畢后需要提供固定的\(\bar\mu,\bar\sigma\)供測試時使用。論文中對所有的mini-batch的\(\mu_\mathcal B,\sigma^2_\mathcal B\)取了均值(m是mini-batch的大小,\(\bar\sigma^2\)采用的是無偏估計)：

\[\begin{align} \bar\mu &=\mathbb E[\mu_\mathcal B] \\ \bar\sigma^2 &={m\over m-1}\mathbb E[\sigma^2_\mathcal B] \end{align} \]

每個mini-batch求\(\bar\mu,\bar\sigma\)這兩個統計量時是對所有的特征點一起計算求的均值和方差.

測試階段，同樣要進行歸一化和縮放平移操作，唯一不同之處是不計算均值和方差，而使用訓練階段記錄下來的\(\bar\mu,\bar\sigma\)。

\[\begin{align} y&=γ({x-\bar\mu\over\sqrt{\bar\sigma^2+ϵ}})+ β \\ &={γ\over\sqrt{\bar\sigma^2+ϵ}}\cdot x+(β-{γ\bar\mu\over\sqrt{\bar\sigma^2+ϵ}}) \end{align} \]

一句話描述 Batch Norm: Batch Norm 層是對每層數據歸一化后再進行線性變換改善數據分布, 其中的線性變換是可學習的.

Batch Norm優點

• 減輕過擬合
• 改善梯度傳播（權重不會過高或過低）
• 容許較高的學習率，能夠提高訓練速度。
• 減輕對初始化權重的強依賴，使得數據分布在激活函數的非飽和區域，一定程度上解決梯度消失問題。
• 作為一種正則化的方式，在某種程度上減少對dropout的使用。

Batch Norm層擺放位置

在激活層（如 ReLU ）之前還是之后，沒有一個統一的定論。在原論文中提出在非線性層之前（CONV_BN_RELU），而在實際編程中很多人可能放在激活層之后(BN_CONV_RELU)。

Batch Norm 應用

Batch Norm在卷積層的應用

前邊提到的mini-batch說的是神經元的個數，而卷積層中是堆疊的多個特征圖，共享卷積參數。如果每個神經元使用一對\(\gamma , \beta\)參數，那么不僅多，而且冗余。可以在channel方向上取m個特征圖作為mini-batch,對每一個特征圖計算一對參數。這樣減少了參數的數量。

應用舉例-VGG16

為VGG16結構模型添加Batch Normalization。

1. 重新完全訓練.如果想將BN添加到卷基層，通常要重新訓練整個模型，大概花費一周時間。
2. finetune.只將BN添加到最后的幾層全連接層，這樣可以在訓練好的VGG16模型上進行微調。采用ImageNet的全部或部分數據按batch計算均值和方差作為BN的初始\(\beta,\gamma\)參數。

與 Dropout 合作

Batch Norm的提出使得dropout的使用減少,但是Batch Norm不能完全取代dropout,保留較小的dropout率,如0.2可能效果更佳。

Batch Norm 實現

caffe的BatchNorm層參數設置示例:

layer {
    name: "conv1/bn"
    type: "BatchNorm"
    bottom: "conv1"
    top: "conv1"
    param { lr_mult: 0 decay_mult: 0 } # mean
    param { lr_mult: 0 decay_mult: 0 } # var
    param { lr_mult: 0 decay_mult: 0 } # scale
    batch_norm_param { use_global_stats: true } # 訓練時設置為 false
}

caffe框架中 BN 層全局均值和方差的實現

與論文計算 global 均值和 global 方差的方式不同之處在於，Caffe 中的 global 均值和 global 方差采用的是滑動衰減平均的更新方式，設滑動衰減系數moving_average_fraction 為 λ，當前的 mini-batch 的均值和方差分別為 \(\mu_B,\sigma_B^2\)：

\[\begin{align} \mu_{new} &= (\lambda \mu_{old} + \mu_B) / s \\ \sigma_{new}^2 &= \begin{cases} (\lambda \sigma_{old}^2 + \frac{m - 1}{m} \sigma_B^2)/s & m > 1 \\ (\lambda \sigma_{old}^2 + \sigma_B^2)/s & m = 1 \end{cases} \\ s &= \lambda s+1 \end{align} \]

簡化形式表示為:$ S_t = (1-\lambda)Y_t + \lambda \cdot S_{t-1} $.
式子中存在一個縮放因子 s 代替BN的batch size, s 初始化為 0, 未采用求所有樣本batch在每一層的平均均值和無偏估計方差的原因是計算不便,需要節約內存和計算資源.在何凱明的caffe實現中僅給出了deploy.prototxt文件方便測試和finetuning.在deploy.prototxt中batch-norm層的參數固定了,均值和方差是在大量數據上嚴格按照論文中的average方法而不是caffe實現中的moving average方法,且數值比較穩定.
caffe中的batch_norm_layer僅含均值方差,不包括gamma/beta,需要后邊緊跟scale_layer,且使用bias對應beta.scale層用於自動學習縮放參數,同時包括了bias_layer能夠學習bias.

官方caffe的batch_norm_layer.cpp代碼如下:

// scale初始化代碼: 用三個blob記錄BatchNorm層的三個數據
void BatchNormLayer<Dtype>::LayerSetUp(...) {
    vector<int> sz;
    sz.push_back(channels_);
    this->blobs_[0].reset(new Blob<Dtype>(sz)); // mean
    this->blobs_[1].reset(new Blob<Dtype>(sz)); // variance
    // 在caffe實現中計算均值方差采用了滑動衰減方式, 用了scale_factor代替num_bn_samples(scale_factor初始為1, 以s=λs + 1遞增).
    sz[0] = 1;
    this->blobs_[2].reset(new Blob<Dtype>(sz)); // normalization factor (for moving average)
}

if (use_global_stats_) {
    // use the stored mean/variance estimates.
    const Dtype scale_factor = this->blobs_[2]->cpu_data()[0] == 0 ?
        0 : 1 / this->blobs_[2]->cpu_data()[0];
    caffe_cpu_scale(variance_.count(), scale_factor,
        this->blobs_[0]->cpu_data(), mean_.mutable_cpu_data());
    caffe_cpu_scale(variance_.count(), scale_factor,
        this->blobs_[1]->cpu_data(), variance_.mutable_cpu_data());
} else {
    // compute mean
    caffe_cpu_gemv<Dtype>(CblasNoTrans, channels_ * num, spatial_dim,
        1. / (num * spatial_dim), bottom_data,
        spatial_sum_multiplier_.cpu_data(), 0.,
        num_by_chans_.mutable_cpu_data());
    caffe_cpu_gemv<Dtype>(CblasTrans, num, channels_, 1.,
        num_by_chans_.cpu_data(), batch_sum_multiplier_.cpu_data(), 0.,
        mean_.mutable_cpu_data());
  }

  // subtract mean
  caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, num, channels_, 1, 1,
      batch_sum_multiplier_.cpu_data(), mean_.cpu_data(), 0.,
      num_by_chans_.mutable_cpu_data());
  caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, channels_ * num,
      spatial_dim, 1, -1, num_by_chans_.cpu_data(),
      spatial_sum_multiplier_.cpu_data(), 1., top_data);

  if (!use_global_stats_) {
    // compute variance using var(X) = E((X-EX)^2)
    caffe_powx(top[0]->count(), top_data, Dtype(2),
        temp_.mutable_cpu_data());  // (X-EX)^2
    caffe_cpu_gemv<Dtype>(CblasNoTrans, channels_ * num, spatial_dim,
        1. / (num * spatial_dim), temp_.cpu_data(),
        spatial_sum_multiplier_.cpu_data(), 0.,
        num_by_chans_.mutable_cpu_data());
    caffe_cpu_gemv<Dtype>(CblasTrans, num, channels_, 1.,
        num_by_chans_.cpu_data(), batch_sum_multiplier_.cpu_data(), 0.,
        variance_.mutable_cpu_data());  // E((X_EX)^2)

    // compute and save moving average
    this->blobs_[2]->mutable_cpu_data()[0] *= moving_average_fraction_;
    this->blobs_[2]->mutable_cpu_data()[0] += 1;
    caffe_cpu_axpby(mean_.count(), Dtype(1), mean_.cpu_data(),
        moving_average_fraction_, this->blobs_[0]->mutable_cpu_data());
    int m = bottom[0]->count()/channels_;
    Dtype bias_correction_factor = m > 1 ? Dtype(m)/(m-1) : 1;
    caffe_cpu_axpby(variance_.count(), bias_correction_factor,
        variance_.cpu_data(), moving_average_fraction_,
        this->blobs_[1]->mutable_cpu_data());
  }

  // normalize variance
  caffe_add_scalar(variance_.count(), eps_, variance_.mutable_cpu_data());
  caffe_powx(variance_.count(), variance_.cpu_data(), Dtype(0.5),
             variance_.mutable_cpu_data());

  // replicate variance to input size
  caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, num, channels_, 1, 1,
      batch_sum_multiplier_.cpu_data(), variance_.cpu_data(), 0.,
      num_by_chans_.mutable_cpu_data());
  caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, channels_ * num,
      spatial_dim, 1, 1., num_by_chans_.cpu_data(),
      spatial_sum_multiplier_.cpu_data(), 0., temp_.mutable_cpu_data());
  caffe_div(temp_.count(), top_data, temp_.cpu_data(), top_data);
  
  caffe_copy(x_norm_.count(), top_data,
      x_norm_.mutable_cpu_data());

BN有合並式和分離式，各有優劣。^[1]

分離式寫法，在切換層傳播時，OS需要執行多個函數，在底層（比如棧）調度上會浪費一點時間。Caffe master branch當前采用的是分離式寫法，CONV層扔掉bias，接一個BN層，再接一個帶bias的SCALE層。
從執行速度來看，合並式寫法需要多算一步bias，參考這里的合並式寫法。

BatchNorm層合並(Conv+BN+Scale+ReLU => Conv+ReLU)

內存優化
在訓練時可以將bn層合並到scale層,在測試(inference)時可以把bn和scale合並到conv層. 在訓練時也可以將凍結的BN層合並到凍結的conv層, 但是不能訓練合並后的conv層, 否則會破壞bn的參數. 合並 bn 層同時可以減少一點計算量.

僅合並BN+Scale=>Scale
bn layer: bn_mean, bn_variance, num_bn_samples 注意在caffe實現中計算均值方差采用了滑動衰減方式,用了scale_factor代替num_bn_samples(scale_factor初始為1,以s=λs+1遞增).
scale layer: scale_weight, scale_bias 代表gamma,beta
BN層的batch的均值mu=bn_mean/num_bn_samples,方差var=bn_variance / num_bn_samples.
scale層設置新的仿射變換參數:

new_gamma = gamma / (np.power(var, 0.5) + 1e-5)
new_beta = beta - gamma * mu / (np.power(var, 0.5) + 1e-5)

Conv+BN+Scale=>Conv
conv layer: conv_weight, conv_bias
在使用BatchNorm時conv_bias通常為0
定義alpha向量為每個卷積核的縮放倍數(長度為通道數),也是特征的均值和方差的縮放因子.

alpha = scale_weight / sqrt(bn_variance / num_bn_samples + eps)
conv_bias = conv_bias * alpha + (scale_bias - (bn_mean / num_bn_samples) * alpha)
for i in range(len(alpha)): conv_weight[i] = conv_weight[i] * alpha[i]

Batch Norm 多卡同步

為什么不進行多卡同步?

BatchNorm的實現都是只考慮了single gpu。也就是說BN使用的均值和標准差是單個gpu算的，相當於縮小了mini－batch size。至於為什么這樣實現，1）因為沒有sync的需求，因為對於大多數vision問題，單gpu上的mini-batch已經夠大了，完全不會影響結果。2）影響訓練速度，BN layer通常是在網絡結構里面廣泛使用的，這樣每次都同步一下GPUs，十分影響訓練速度。^[2]

但是為了達到更好的效果, 實現Sync-BN也是很有意義的.

在深度學習平台框架中多數是采用數據並行的方式, 每個GPU卡上的中間數據沒有關聯.
為了實現跨卡同步BN, 在前向運算的時候需要計算全局的均值和方差，在后向運算時候計算全局梯度。 最簡單的實現方法是先同步求均值，再發回各卡然后同步求方差，但是這樣就同步了兩次。實際上均值和方差可以放到一起求解, 只需要同步一次就可以. 數據並行的方式改為下圖所示:^[3]

多卡同步的公式原理^[4]

\[\begin{align} \mu &= \frac{1}{m}\sum_{i=1}^m x_i \\ \sigma^2 &= \frac{1}{m} \sum_{i=1}^m (x_i - \mu)^2 = \frac{1}{m} \sum_{i=1}^m (x_i^2+\mu^2-2x_i\mu) = \frac{1}{m} \sum_{i=1}^m x_i^2 - \mu^2 \\ &= \frac{1}{m} \sum_{i=1}^m x_i^2 - (\frac{1}{m}\sum_{i=1}^m x_i)^2 \end{align} \]

因此總體batch_size對應的均值和方差可以通過每張GPU中計算得到的 \(\sum x_i\) 和 \(\sum x_i^2\) reduce相加得到. 在反向傳播時也一樣需要同步一次梯度信息.

另外, 可以參考sync-bn 實現討論.

---

1. 從Bayesian角度淺析Batch Normalization ↩︎

2. https://www.zhihu.com/question/59321480/answer/198077371 ↩︎

3. MXNet Gluon 上實現跨卡同步 Batch Normalization ↩︎

4. Context Encoding for Semantic Segmentation ↩︎