为什么 Keras BatchNorm 产生的输出与 PyTorch 不同?

Why does Keras BatchNorm produce different output than PyTorch?

手电筒:'1.9.0+cu111'

Tensorflow-GPU:'2.5.0'

遇到一个奇怪的事情,用tensorflow 2.5的Batch Normal层和Pytorch 1.9的BatchNorm2d层计算同一个Tensor,结果完全不同 (TensorFlow接近1,Pytorch接近0)一开始以为是momentum和epsilon的区别,改成一样后,结果一样

from torch import nn
import torch
x = torch.ones((20, 100, 35, 45))
a = nn.Sequential(
            # nn.Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), padding=0, bias=True),
            nn.BatchNorm2d(100)
        )
b = a(x)

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras.layers import *
x = tf.ones((20, 35, 45, 100))
a = keras.models.Sequential([
            # Conv2D(128, (1, 1), (1, 1), padding='same', use_bias=True),
            BatchNormalization()
        ])
b = a(x)

批量归一化在训练和推理中的工作方式不同,

在训练期间(即当使用fit()或使用参数training=True调用layer/model时),该层对其输出进行归一化使用当前批次输入的均值和标准差。也就是说,对于每个被归一化的通道,returns

gamma * (batch - mean(batch)) / sqrt(var(batch) + epsilon) + beta

其中:

  • epsilon 是小常量(可配置为构造函数参数的一部分)
  • gamma 是学习的缩放因子(初始化为 1),可以通过将 scale=False 传递给构造函数来禁用它。
  • beta是学习的偏移因子(初始化为0),可以通过将center=False传递给构造函数来禁用它。

在推理期间(即当使用 evaluate()predict() 或使用参数 training=False 调用 layer/model 时(这是默认值),该层使用它在训练期间看到的批次的均值和标准差的移动平均值对其输出进行归一化。也就是说,它 returns

gamma * (batch - self.moving_mean) / sqrt(self.moving_var + epsilon) + beta.

self.moving_meanself.moving_var 是不可训练的变量,每次在训练模式下调用层时都会更新,例如:

    moving_mean = moving_mean * momentum + mean(batch) * (1 - momentum)
    moving_var = moving_var * momentum + var(batch) * (1 - momentum)

参考:https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization

如果你在eval模式下运行 pytorch batchnorm,你会得到接近的结果(其余差异来自不同的内部实现、参数选择等),

from torch import nn
import torch
x = torch.ones((1, 2, 2, 2))
a = nn.Sequential(
            # nn.Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), padding=0, bias=True),
            nn.BatchNorm2d(2)
        )
a.eval()
b = a(x)
print(b)
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras.layers import *
x = tf.ones((1, 2, 2, 2))
a = keras.models.Sequential([
            # Conv2D(128, (1, 1), (1, 1), padding='same', use_bias=True),
            BatchNormalization()
        ])
b = a(x)
print(b)

out:

tensor([[[[1.0000, 1.0000],
          [1.0000, 1.0000]],

         [[1.0000, 1.0000],
          [1.0000, 1.0000]]]], grad_fn=<NativeBatchNormBackward>)
tf.Tensor(
[[[[0.9995004 0.9995004]
   [0.9995004 0.9995004]]

  [[0.9995004 0.9995004]
   [0.9995004 0.9995004]]]], shape=(1, 2, 2, 2), dtype=float32)