前記

最近感覺Batch Normalization在實踐上還是很有搞頭的，首先在單機TensorFlow里完整實現BN演算法就有很多坑（具體可以看我的文章Michael：[論文閱讀]Batch Normalization: Accelerating Deep Netwo，裡面也有我對BN的理解），然後在單機多卡上的實現也有些坑，主要是moving mean和moving variance在單機多卡上是怎麼實現更新的（分散式的原理類似於單機多卡）。

首先做一些實驗來看看BN的運行原理吧，代碼已上傳至testBNInMultiGPU.py，我改了一些bug後可以完美運行並在實際項目中運行成功的代碼在validSyncBN.py。

單機單卡上的BN實驗

代碼為testBNInSingleGPU()。

結果：

outputArray = [0. 0. 0.]

After normalization:
moving mean = [0.1 0.1 0.1]
moving variance = [0.9 0.9 0.9]

首先看輸出全為0表示BN起到了歸一化的作用。

之後可以看moving mean和moving variance的值，參考Michael：TensorFlow源碼解讀之Batch Normalization，得知TF中moving mean和moving variance的更新方式都是momentum，即 x=m imes x+(1-m) imes E_B(x)=x+(1-m) imes(E_B(x)-x) 。

因為moving mean初始化為0，moving variance初始化為1，momentum設置為0.9，代入可得：

moving mean： 0+(1-0.9) imes(1-0)=0.1

moving variance： 1+(1-0.9) imes(0-1)=0.9

和實驗是完美符合的。

單機多卡上的運行實驗

這裡選取了2塊卡進行實驗，代碼為testBNInMultiGPU()。

結果：

outputArray = [0. 0. 0.]

After normalization:
moving mean = [0.2 0.2 0.2]
moving variance = [0.79999995 0.79999995 0.79999995]

由於是兩塊卡，結合Michael：TensorFlow源碼解讀之Batch Normalization的分析，再根據實驗結果，我推測更新公式是這樣的：

統計每張卡上moving mean和moving variance的變化量，即 Delta x=x-x=(1-m) imes(E_B(x)-x) ，將各卡的變化量求和得到總變化量，進而得到更新的結果：

x=x+sum^N_1{Delta x}=x+sum^N_1{(1-m) imes(E_B(x)-x)} ，這裡 N 表示GPU的數量。

因為moving mean初始化為0，moving variance初始化為1，momentum設置為0.9，代入可得：

moving mean： 0+(1-0.9) imes(1-0)+(1-0.9) imes(1-0)=0.2

moving variance： 1+(1-0.9) imes(0-1)+(1-0.9) imes(0-1)=0.8

和實驗是完美符合的。

另一種更新的策略是，每張卡前向執行一次，更新一次moving mean和moving variance，那麼：

moving mean：

第一次： 0+(1-0.9) imes(1-0)=0.1

第二次： 0.1+(1-0.9) imes(1-0.1)=0.19

這與實驗結果是不符合的。

這裡改變一下輸入的值（一個卡上是1，一個卡上是2），以驗證我的想法是不是對的：

for item in range(numberGPU):
feedDict[tf.get_default_graph().get_tensor_by_name("tower_%s/data:0" % item)]
= np.ones((3))*(item+1)

結果：

outputArray = [0. 0. 0.]

After normalization:
moving mean = [0.3 0.3 0.3]
moving variance = [0.79999995 0.79999995 0.79999995]

因為moving mean初始化為0，moving variance初始化為1，momentum設置為0.9，代入可得：

moving mean： 0+(1-0.9) imes(1-0)+(1-0.9) imes(2-0)=0.3

moving variance： 1+(1-0.9) imes(0-1)1+(1-0.9) imes(0-1)=0.8

和實驗是完美符合的。

影響

對於公式 x=x+sum^N_1{Delta x}=x+sum^N_1{(1-m) imes(E_B(x)-x)} ，如果各個卡上 x 的分布是一致的，那麼很清晰就得到：

Delta x=N imes (1-m) imes(E_B(x)-x)

這個實際是會加速moving mean和moving variance的，所以可設置新的moment為： m=frac{N+m-1}{N} 。

第二個影響就是對訓練了，從上述的實驗可以猜測BN的訓練過程：

在每張卡上前向推導，得到BN的moving mean和moving variance的改變數，同時，反向推導，得到 gamma 和 eta 在每張卡上的梯度。

然後，使用tensorflow 多GPU編程 完全指南中描述的方法，得到gamma 和 eta的平均梯度。

所以，TensorFlow的官方方法是沒有考慮多GPU編程的，下一個步驟好好搞搞這個吧，目測比較難。

目前找到的一些比較好的資源：

batch normalization的multi-GPU版本該怎麼實現？

裡面有幾個比較好的解決方案：

1.jianlong-yuan/syncbn-tensorflow重寫了TensorFlow的官方方法，可以做個實驗驗證一下。

2.曠視科技：CVPR 2018 | 曠視科技物體檢測冠軍論文——大型Mini-Batch檢測器MegDet

tensorpack/tensorpack裡面有該論文的代碼實現。

關於第一種解決方案的實驗論證與理論解釋

jianlong-yuan/syncbn-tensorflow是源代碼，我稍微改了改代碼，並添加了我的一些注釋，代碼為sync_batch_norm()。

測試代碼為testBNInMultiGPU3()。

輸出為：

outputTuple = [array([-0.998006, -0.998006, -0.998006], dtype=float32), array([0.9980061, 0.9980061, 0.9980061], dtype=float32),
[array([0.15, 0.15, 0.15], dtype=float32), array([0.92499995, 0.92499995, 0.92499995], dtype=float32)]]

After normalization:
moving mean = [0.15 0.15 0.15]
moving variance = [0.92499995 0.92499995 0.92499995]

outputTuple一共4項，前兩項為2個GPU的輸出量，一個為-1，一個為1，後兩項分別是moving mean和moving variance的輸出。

這裡先就數據的觀點來驗證結果是否正確。

GPU_0的數據是3個1，GPU_1的數據是3個2，參考batch normalization accelerating deep network training by reducing internal covariate shift演算法1，有：

mu_B=frac{3 imes1+3 imes2}{6}=1.5

sigma^2_B=frac{3 imes(1-1.5)^2+3 imes(2-1.5)^2}{6}=0.25

對於GPU_0的數據，有：

hat x=frac{x-mu_B}{sqrt{sigma^2_B+epsilon}}=frac{1-1.5}{sqrt{0.25}}=-1

y=gamma hat x+eta=1 imes-1+0=-1

對於GPU_1的數據，有：

hat x=frac{x-mu_B}{sqrt{sigma^2_B+epsilon}}=frac{2-1.5}{sqrt{0.25}}=1

y=gamma hat x+eta=1 imes1+0=1

這說明輸出是符合的。

下面看moving mean和moving variance的值。

moving mean： 0+(1-0.9) imes(1.5-0)=0.15

moving variance： 1+(1-0.9) imes(0.25-1)=0.925

也是符合的，說明程序還是很魯棒的。

這裡來分析程序的運行原理，大部分我都用注釋說明了，這裡講裡面的兩個知識點吧，一個是mean和variance是怎麼計算得到的，一個是moving mean和moving variance是怎麼更新的。

首先說下mean和variance是怎麼計算得到的，是以下的這段代碼：

# avarage moving_mean and moving_var in multi GPUs
shared_name = tf.get_variable_scope().name
batch_mean = tf.reduce_mean(inputs, axis=axis)
batch_mean = gen_nccl_ops.nccl_all_reduce(
input=batch_mean,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean) * (1.0 / num_dev)
batch_mean_square = tf.reduce_mean(tf.square(inputs), axis=axis)
batch_mean_square = gen_nccl_ops.nccl_all_reduce(
input=batch_mean_square,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean_square) * (1.0 / num_dev)
mean = batch_mean
var = batch_mean_square - tf.square(batch_mean)

假設有N塊卡，每個卡上有 m_i,i=1, 2, ...,N 個訓練樣本，且令 m_1=m_2=...=m_N ，那麼有 m=m_1*N=m_2*N=...=m_N*N ，其中， m 是整個的樣本數。

這句代碼：

batch_mean = tf.reduce_mean(inputs, axis=axis)

其用數學表達式可以表示成： batch\_mean_i=frac{sum_{j=1}^{m_i}x_j}{m_i},for,i=1,2, ...,N

這句代碼：

batch_mean = gen_nccl_ops.nccl_all_reduce(
input=batch_mean,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean) * (1.0 / num_dev)

其會在各卡對相應的變數進行操作，這裡進行的操作是求和，最後除N即可，用數學表達式可以表達成：

batch\_mean=frac{sum_i^Nbatch\_mean_i}{N}=frac{sum_i^Nfrac{sum_{j=1}^{m_i}x_j}{m_i}}{N} ，因為有 m_1=m_2=...=m_N ，那麼有：

batch\_mean=frac{sum_i^N sum_{j=1}^{m_i}x_j}{Nm_1}=frac{sum_{j=1}^mx_j}{m}

這個表達式就與論文里的表達式一致了，以下用 mu_B 代替 batch\_mean 。

這句代碼：

batch_mean_square = tf.reduce_mean(tf.square(inputs), axis=axis)

其用數學表達式可以表示成： batch\_mean\_square_i=frac{sum_{j=1}^{m_i}x_j^2}{m_i},for , i=1,2,...,N

這句代碼：

batch_mean_square = gen_nccl_ops.nccl_all_reduce(
input=batch_mean_square,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean_square) * (1.0 / num_dev)

其用數學表達式可以表示成： batch\_mean\_square=frac{sum_i^Nbatch\_mean\_square_i}{N}=frac{sum_i^Nfrac{sum_{j=1}^{m_i}x_j^2}{m_i}}{N}\ =frac{sum_i^Nsum_{j=1}^{m_i}x_j^2}{N m_1}=frac{sum_{j=1}^mx_j^2}{m}

這句代碼：

var = batch_mean_square - tf.square(batch_mean)

其用數學表達式可以表示成： var=batch\_mean\_square-mu_B^2=frac{sum_{j=1}^mx_j^2}{m}-mu_B^2=frac{sum_{j=1}^mx_j^2-mmu_B^2}{m}

這裡利用的一個trick是 mu_B=frac{sum_j^mx_j}{m} ，那麼 var=frac{sum_{i=1}^m(x_i-mu_B)^2}{m}=frac{sum_{i=1}^mx_i^2-2mu_Bsum_{i=1}^mx_i+mu_B^2}{m}=frac{sum_{i=1}^mx_i^2-mu_B^2}{m}

另外一個值得注意的代碼就是這個控制條件：

if int(outputs.device[-1]) == 0:

因為每個GPU上的moving mean和moving variance都一致了，所以沒必要每個都更新，更新GPU:0的即可。

綜上，看起來這個代碼還是不錯的，待我在實際項目中再驗證驗證吧。

在實際項目中的實驗論證

說實話我是很想在實驗中加入進去的，可惜的是，在我本地跑沒有問題，但是在伺服器上跑最終依然有問題，所以最後可能得放棄這個了。

最開始碰到的問題在於Python 2中的對str的解釋，報錯為：

Traceback (most recent call last):
File "/home/dao/anaconda2/envs/python36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2349, in get_attr
c_api.TF_OperationGetAttrValueProto(self._c_op, name, buf)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Operation tower_0/lstm2fc/batch_normalization_0/NcclAllReduce_1 has no attr named _XlaCompile.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/home/dao/anaconda2/envs/python36/lib/python3.6/site-packages/tensorflow/python/ops/gradients_impl.py", line 391, in _MaybeCompile
xla_compile = op.get_attr("_XlaCompile")
File "/home/dao/anaconda2/envs/python36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2353, in get_attr
raise ValueError(str(e))
ValueError: Operation tower_0/lstm2fc/batch_normalization_0/NcclAllReduce_1 has no attr named _XlaCompile.
...

在搜索解決方案的時候我碰巧發現本文所寫的第二種解決方案（即曠視論文）與第一種解決方案異曲同工（沒興趣去了解代碼的根源了）。

第二種解決方案的代碼在tensorpack/tensorpack/blob/master/tensorpack/models/batch_norm.py，首先給了我一些關於TF支持nccl的解釋：

assert six.PY2 or TF_version >= (1, 10),
"Cross-GPU BatchNorm is only supported in TF>=1.10 ."
"Upgrade TF or apply this patch manually: https://github.com/tensorflow/tensorflow/pull/20360"

鏈接的issue是Fix gradient of nccl_ops by ppwwyyxx · Pull Request #20360 · tensorflow/tensorflow，個人猜測1.10以下的版本支持Python 2.7的居多，我用Python 3.6，所以報錯了。

升級到了1.10.0的版本，發現代碼這段加上也是很不錯的：

if TF_version <= (1, 12):
try:
from tensorflow.contrib.nccl.python.ops.nccl_ops import _validate_and_load_nccl_so
except Exception:
pass
else:
_validate_and_load_nccl_so()
from tensorflow.contrib.nccl.ops import gen_nccl_ops
else:
from tensorflow.python.ops import gen_nccl_ops

弄好之後又報錯：

tensorflow.python.framework.errors_impl.NotFoundError: libnccl.so.2

參考文章tensorflow1.12 多GPU協同訓練報錯tensorflow.python.framework.errors_impl.NotFoundError: libnccl.so.2發現需要安裝nccl庫，Nvidia官方的安裝指導在Deep Learning SDK Documentation。

好歹不報錯了，但是程序似乎不運行了，也查不出原因，所以決定暫時放棄這種做法了。

後續更新：排除bug

終於搞通了，發現是代碼的bug，Tensorflow的條件控制沒有if else這麼簡單，必須寫到計算圖裡面去。

這會導致一個什麼結果呢？在推理階段，不需要更新moving mean和moving variance，所以代碼會直接推理得到結果：

outputs, _, _ = tf.nn.fused_batch_norm(inputs, gamma, beta, mean=moving_mean, variance=moving_var, epsilon=epsilon, is_training=False)

但是呢，這是一個else語句，計算圖是沒法把Python里的判斷條件放到計算圖裡去的，如果列印計算圖，會發現此時根本就沒有直接推理的結果：

print ([n.name for n in tf.get_default_graph().as_graph_def().node if "fused_batch_norm" in n.name])

使用tf.cond()寫好了判斷函數後，還發現一個問題，在我真實的項目環境中是可以使用的，但是我在寫的測試代碼中卻會報這樣的錯：

"operation has been marked as not fetchable"

參考Queue operation in conditional execution context fails with "operation has been marked as not fetchable" · Issue #4094 · tensorflow/tensorflow，控制條件不能有太多控制語句，我發現True分支里寫了moving mean和moving variance的更新語句，這個會在分支里引入過多的輸入，我猜測TensorFlow對此做了限定，在實驗中我不在sess.run()中求tf.get_collection(tf.GraphKeys.UPDATE_OPS)的值也是可以不會報錯的，綜上，我認為錯誤在於更新語句，也做了一些改進，詳細代碼見validSyncBN.py。

【已完結】


推薦閱讀：
相关文章