【導讀】本文是谷歌機器學習工程師 Chris Rawles 撰寫的一篇技術博文，探討了如何在 TensorFlow 和 tf.keras 上利用 Batch Normalization 加快深度神經網路的訓練。我們知道，深度神經網路一般非常複雜，即使是在當前高效能GPU的加持下，要想快速訓練深度神經網路依然不容易。Batch Normalization 也許是一個不錯的加速方法，本文介紹了它如何幫助解決梯度消失和梯度爆炸問題，並討論了ReLu啟用以及其他啟用函式對於抵消梯度消失問題的作用。最後，本文使用TensorFlow和tf.keras實現了在MNIST上Batch Normalization，有助於加深讀者理解。
How to use Batch Normalization with TensorFlow and tf.keras to train deep neural networks faster
訓練深度神經網路可能非常耗時。但是可以通過消除梯度來顯著地減少訓練時間，這種情況發生在網路由於梯度（特別是在較早的層中的梯度）接近零值而停止更新。 結合Xavier權重初始化和ReLu啟用功能有助於抵消消失梯度問題。 這些技術也有助於解決與之相反的梯度爆炸問題，這種情況下梯度變得非常大，它防止模型更新。
批量標準化（Batch Normalization）也許是對付梯度消失和爆炸問題的最有力工具。 批量標準化的工作方式如下：對於給定層中的每個單元，首先計算z分數，然後在兩個受過訓練的變數γ和β應用線性轉換。 批量標準化通常在非線性啟用函式之前完成（參見下文），但在啟用函式之後應用批量標準也可能是有利的。 檢視這個講座瞭解該技術如何工作的更多細節。
http://cs231n.stanford.edu/slides/2017/cs231n_2017_lecture6.pdf
在反向傳播過程中，梯度傾向於在較低層裡變得更小，從而減緩權重更新並因此減少訓練次數。 批量標準化有助於消除所謂的梯度消失問題。
批量標準化可以在TensorFlow中以三種方式實現。使用：
1. tf.keras.layers.BatchNormalization
2. tf.layers.batch_normalization
3. tf.nn.batch_normalization
tf.keras模組成為1.4版TensorFlow API的核心的一部分。 併為構建TensorFlow模型提供高階API; 所以我會告訴你如何在Keras做到這一點。tf.layers.batch_normalization函式具有類似的功能，但Keras被證明是在TensorFlow中編寫模型函式的一種更簡單的方法。
in_training_mode = tf.placeholder(tf.bool)
hidden = tf.keras.layers.Dense(n_units,
activation=None)(X) # no activation function, yet
batch_normed = tf.keras.layers.BatchNormalization()(hidden, training=in_training_mode)
output = tf.keras.activations
.relu(batch_normed) # ReLu is typically done after batch normalization
# optimizer code here …
extra_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(extra_ops):
train_op = optimizer.minimize(loss)
注意批量標準化函式中的訓練變數。 這是必需的，因為批量標準化在訓練期間與應用階段的操作方式不同。在訓練期間，z分數是使用批均值和方差計算的，而在推斷中，則是使用從整個訓練集估算的均值和方差計算的。
在TensorFlow中，批量標準化可以使用tf.keras.layers作為附加層實現。
包含tf.GraphKeys.UPDATE_OPS的第二個程式碼塊很重要。對於網路中的每個單元，使用tf.keras.layers.BatchNormalization，TensorFlow會不斷估計訓練資料集上權重的均值和方差。這些儲存的值用於在預測時間應用批量標準化。 每個單元的訓練集均值和方差可以通過列印extra_ops來觀察，extra_ops包含網路中每圖層的列表：
print(extra_ops)
[<tf.Tensor ‘batch_normalization/AssignMovingAvg:0’ shape=(500,) dtype=float32_ref>,
# layer 1 mean values
<tf.Tensor ‘batch_normalization/AssignMovingAvg_1:0’ shape=(500,) dtype=float32_ref>,
# layer 1 variances ...]
雖然批量標準化在tf.nn模組中也可用，但它需要額外的記錄，因為均值和方差是函式的必需引數。 因此，使用者必須在批次級別和訓練集級別上手動計算均值和方差。 因此，它是一個比tf.keras.layers或tf.layers更低的抽象層次;應避免用tf.nn實現。
▌在MNIST上批量標準化
下面，我使用TensorFlow將批量標準化應用到突出的MNIST資料集。 看看這裡的程式碼。 MNIST是一個易於分析的資料集，不需要很多層就可以實現較低的分類錯誤。 但是，我們仍然可以構建深度網路並觀察批量標準化如何實現收斂。
我們使用tf.estimator API構建自定義估算器。 首先我們建立模型：
def dnn_custom_estimator(features, labels, mode, params):
in_training = mode == tf.estimator.ModeKeys.TRAIN
use_batch_norm = params['batch_norm']
net = tf.feature_column.input_layer(features, params['features'])
for i, n_units in enumerate(params['hidden_units']):
net = build_fully_connected(net, n_units=n_units, training=in_training,
batch_normalization=use_batch_norm,
activation=params['activation'],
name='hidden_layer' + str(i))
logits = output_layer(net, 10, batch_normalization=use_batch_norm,
training=in_training)
predicted_classes = tf.argmax(logits, 1)
loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)
accuracy = tf.metrics.accuracy(labels=tf.argmax(labels, 1),
predictions=predicted_classes,
name='acc_op')
tf.summary.scalar('accuracy', accuracy[1]) # for visualizing in TensorBoard
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode, loss=loss,
eval_metric_ops={'accuracy': accuracy})
# Create training op.
assert mode == tf.estimator.ModeKeys.TRAIN
extra_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
optimizer = tf.train.AdamOptimizer(learning_rate=params['learning_rate'])
with tf.control_dependencies(extra_ops):
train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
在我們定義模型函式之後，讓我們構建自定義估計器並訓練和評估我們的模型：
def train_and_evaluate(output_dir):
features = [tf.feature_column.numeric_column(key='image_data', shape=(28*28))]
classifier = tf.estimator.Estimator(model_fn=dnn_custom_estimator,
model_dir=output_dir,
params={'features': features,
'batch_norm': USE_BATCH_NORMALIZATION,
'activation': ACTIVATION,
'hidden_units': HIDDEN_UNITS,
'learning_rate': LEARNING_RATE})
train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=NUM_STEPS)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn)
tf.estimator.train_and_evaluate(classifier, train_spec, eval_spec)
train_and_evaluate('mnist_model')
讓我們來測試批量標準化如何影響不同深度的模型。將我們的程式碼打包到Python包後，我們可以使用Cloud ML Engine並行執行多個實驗：
# def ml-engine function
submitMLEngineJob() {
gcloud ml-engine jobs submit training $JOBNAME
--package-path=$(pwd)/mnist_classifier/trainer
--module-name trainer.task
--region $REGION
--staging-bucket=gs://$BUCKET
--scale-tier=BASIC
--runtime-version=1.4
--
--outdir $OUTDIR
--hidden_units $net
--num_steps 1000
$batchNorm
}
# launch jobs in parallel
export PYTHONPATH=${PYTHONPATH}:${PWD}/mnist_classifier
for batchNorm in '' '--use_batch_normalization'
do
net=''
for layer in 500 400 300 200 100 50 25;
do
net=$net$layer
netname=${net//,/_}${batchNorm/--use_batch_normalization/_bn}
JOBNAME=mnist$netname_$(date -u +%y%m%d_%H%M%S)
OUTDIR=gs://${BUCKET}/mnist_models/mnist_model$netname/trained_model
echo $OUTDIR $REGION $JOBNAME
gsutil -m rm -rf $OUTDIR
submitMLEngineJob
net=$net,
done
done
下圖顯示了達到90％測試精度所需的訓練迭代次數（1次迭代包含的批次大小為500）。 很明顯，批量標準化顯著加快了深度網路的訓練。如果沒有批量標準化，隨著每個後續層的增加，訓練步驟的數量都會增加，但使用它後，訓練步數幾乎保持不變。 在實踐中，它是面對更困難的資料集，更多層網路結構時取得成功的先決條件。
如果沒有批量標準化，達到90％準確度所需的訓練迭代次數會隨著層數的增加而增加，這可能是由於梯度消失造成的。
同樣，如下所示，對於具有7個隱藏層的全連線的網路，沒有批量標準化的收斂時間較慢
上述實驗利用了常用的ReLu啟用功能。 雖然不能像上面所示一樣抵擋梯度消失帶來的效應，ReLu啟用比Sigmoid或tanh啟用功能要好得多。 Sigmoid啟用函式對梯度消失很無力。在更大的數值（非常正或負）時，sigmoid函式“飽和” 即S形函式的導數接近零。 當越來越多節點飽和時，更新次數減少，網路停止訓練。
使用sigmoid啟用函式而不使用批量標準化，相同的7層網路訓練會顯著減慢。當使用批量標準化，網路達到收斂時的迭代次數與使用ReLu相似。
另一方面，其他啟用函式（如指數ReLu或洩漏ReLu函式）可以幫助抵制梯度消失問題，因為它們對於正數和負數都具有非零導數。
最後，重要的是要注意批量標準化會給訓練帶來額外的時間成本。 儘管批量標準化通常會減少達到收斂的訓練步數，但它會帶來額外的時間成本，因為它引入了額外的操作，並且還給每個單元引入了兩個新的訓練引數。
對於MNIST分類問題（使用1080 GTX GPU），批量標準化能在較少的迭代次數收斂，但每次迭代的時間較慢。 最終，批量標準化版本的收斂速度仍然較快，但整合訓練時間後，改進效果並不明顯。
結合XLA和混合批量標準化（fused Batch Normalization）（在tf.layers.batch_normalization中融合了引數）可以通過將幾個單獨的操作組合到單個核心中來加速批量標準化操作。
無論如何，批量標準化可以成為加速深度神經網路訓練的非常有價值的工具。 像訓練深度神經網路一樣，確定一種方法是否有助於解決問題的最佳方法就是做一下實驗！
附Keras官方資源
- 官網：keras.io
- 中文版文件：keras.io/zh/
- 快速入門：keras.io/zh/#30-kera…
- Github：github.com/keras-team/…
- Google+網上論壇：groups.google.com/forum/#!for…
- slack：kerasteam.slack.com/
參考文獻：
https://towardsdatascience.com/how-to-use-batch-normalization-with-tensorflow-and-tf-keras-to-train-deep-neural-networks-faster-60ba4d054b73