Pytorch Batchnorm 层与 Keras Batchnorm 不同

Question

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

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

火炬：

平均值和标准差是按小批量计算的

来源： Pytorch BatchNorm

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

凯拉斯：

axis：整数，应该归一化的轴（通常是特征轴）。 例如，在一个带有 data_format="channels_first" 的 Conv2D 层之后，在 BatchNormalization 中设置 axis=1。

来源： Keras BatchNorm

在这里他们对特征（通道）进行平均

什么是正确的方法？ 如何在模型之间传递 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])