python - Pytorch Batchnorm 层与 Keras Batchnorm 不同

Question

我正在尝试将预训练的 BN 权重从 pytorch 模型复制到其等效的 Keras 模型，但我不断得到不同的输出。

我阅读了 Keras 和 Pytorch BN 文档，我认为区别在于它们计算“均值”和“变量”的方式。

torch :

The mean and standard-deviation are calculated per-dimension over the mini-batches

来源:Pytorch BatchNorm

因此，他们对样本进行平均。

凯拉斯:

axis: Integer, the axis that should be normalized (typically the features axis). For instance, after a Conv2D layer with data_format="channels_first", set axis=1 in BatchNormalization.

来源:Keras BatchNorm

在这里他们对特征( channel )进行平均

什么是正确的方法？如何在模型之间传递 BN 权重？

Answer 1

您可以从 pytorch 模块的 running_mean 和 running_var 属性中检索 moving_mean 和 moving_variance

# torch weights, bias, running_mean, running_var corresponds to keras gamma, beta, moving mean, moving average

weights = torch_module.weight.numpy()  
bias = torch_module.bias.numpy()  
running_mean =  torch_module.running_mean.numpy()
running_var =  torch_module.running_var.numpy()

keras_module.set_weights([weights, bias, running_mean, running_var])