推論時のBatch Normalization(バッチ正規化)の各種パラメータの仕様について

Question

前提・実現したいこと

推論時における、BN(Batch Normalization)層での学習で決定されたパラメータがどのように正規化に用いられるのか知りたい。

現在の状況

PyTorchであるネットワークを自分で実装し、パラメータ抽出を行ったがBN層におけるパラメータの仕様がどうもわからないです。
というのも
MobileNet V1内のバッチ正規化(Batch Normalization)層の推論時の入力データに対する処理内容を知りたい
こちらの質問で複数BN層がモデル内にあった場合、各層で１組の平均と標準偏差が決定されると思っていました。
しかし実際に

for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

とすると

Model's state_dict:
features.0.weight 	 torch.Size([32, 3, 3, 3])
features.0.bias 	 torch.Size([32])
features.2.dw_conv_k3.weight 	 torch.Size([32, 1, 3, 3])
features.2.bn.weight 	 torch.Size([32])
features.2.bn.bias 	 torch.Size([32])
features.2.bn.running_mean 	 torch.Size([32])
features.2.bn.running_var 	 torch.Size([32])
features.2.bn.num_batches_tracked 	 torch.Size([])
features.3.p_conv_k1.weight 	 torch.Size([64, 32, 1, 1])
features.3.bn.weight 	 torch.Size([64])
features.3.bn.bias 	 torch.Size([64])
features.3.bn.running_mean 	 torch.Size([64])
features.3.bn.running_var 	 torch.Size([64])
features.3.bn.num_batches_tracked 	 torch.Size([])
....(以下略)

と出力され、第二層目のBN層の学習で決定されたパラメータ(features.2.bn.running_mean等)が32要素の配列としてあります。
(私の予想としては1要素で構成された各層に対するパラメータの組が出力されると思っていました。)
ここで例としてfeatures.2.bn.running_mean[32]の中身をのぞくと

tensor([ 0.2491, -0.0099,  0.0855,  0.1906,  0.0792, -0.3274,  0.1849, -0.4113,
         0.0127,  0.3673,  0.2680,  0.0913,  0.1177,  0.0630, -0.1980,  0.0917,
         0.0147,  0.1684,  0.0583, -0.2606, -0.1515,  0.0202, -0.1987,  0.0284,
        -0.0948, -0.7597, -0.1758,  0.2207, -0.4727,  0.1047, -0.1240,  0.2810],
       device='cuda:0')

となっておりしっかりと配列の中身詰まっているのがわかりました。

聞きたいこと

ここで考えられるのが現在私が実装しているモデルでは第二層目のBN層の入力では32チャネルの入力画像が、第三層目のBN層では64チャネルの入力画像が与えられるため、抽出されたfeatures.xx.bn.running_meanやfeatures.xx.bn.running_varの各要素を各入力チャネルに対応させて推論時に正規化
(つまり32要素のbn.runnig_meanがあった場合、32チャネルの入力画像に対して、1チャネル目にはbn.runnig_meanの1要素目を、2チャネル目には2要素目を適用...といった形)
を行っているということなのですが、この理解で正しいでしょうか？
どなたか助言を頂ければ助かります。

Answer 1

各層で１組の平均と標準偏差が決定されると思っていました。

機械学習で標準化といった場合、普通チャンネル (特徴量) ごとに行うので、前の出力のチャンネル数が32なら、Batch Normalization 層のパラメータが数が全部32なのは別におかしくないと思います。

features.2.bn.weight      torch.Size([32])
features.2.bn.bias      torch.Size([32])
features.2.bn.running_mean      torch.Size([32])
features.2.bn.running_var      torch.Size([32])

features.xx.bn.running_varの各要素を各入力チャネルに対応させて推論時に正規化

(つまり32要素のbn.runnig_meanがあった場合、32チャネルの入力画像に対して、1チャネル目にはbn.runnig_meanの1要素目を、2チャネル目には2要素目を適用...といった形)
を行っているということなのですが、この理解で正しいでしょうか？

その理解であってます。

Comments

[Otom_jp]

ご回答ありがとうございます

[tiitoi]

はい、そうです。