工程Trick | 合併BN層加速前向推理

1.BatchNorm

在訓練深度網絡模型時,BN(Batch Normalization)層能夠加速網絡收斂,並且能夠控制過擬合,一般放在卷積層之後。如下示意圖所示,BatchNorm是以通道爲單位,對當前通道中所有的N、H、W做歸一化。

BN 層將特徵歸一化後,能夠有效解決由於特徵分佈不均勻而導致的梯度消失與梯度爆炸問題。並通過可學習參數保證特徵的有效性。雖然 BN 層在訓練時起到了積極作用,然而,在網絡前向推斷時多了一些層的運算,影響了模型的性能,且佔用了更多的內存或者顯存空間。基於此將卷積後的BN操作直接融合到卷積中,這是現在工程上普遍會使用的方法。

2.將BN操作合併到卷積中

3.pytorch代碼實現

很明顯,下述代碼就是將上部分中的公式進行了代碼化。在這裏不得不說一句pytorch大法NB。

conv_param = model_state_dict[merge_bn_model_dict_key[index]]
conv_bias = model_state_dict[merge_bn_model_dict_key[index+1]]
outchannel = conv_param.size()[1]
bn_weignt = model_state_dict[merge_bn_model_dict_key[index+2]]
bn_bias = model_state_dict[merge_bn_model_dict_key[index+3]]
bn_mean = model_state_dict[merge_bn_model_dict_key[index+4]]
bn_var = model_state_dict[merge_bn_model_dict_key[index+5]]

#merge
var_sqrt = torch.sqrt(bn_var+bn_eps)
merge_bn_conv_param = conv_param*((bn_weight/var_sqrt)
merge_bn_conv_bias = bn_bias + bn_weight*(conv_bias-bn_mean/var_sqrt)

new_conv_key = merge_bn_model_dict_key[index]
new_bias_key = new_conv_key[:,-6]+'bias'


merge_bn_model_dict[new_conv_key] = merge_bn_conv_param
merge_bn_model_dict[new_bias_key] = merge_bn_conv_bias
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章