Tensorflow模型训练正常但推理输出NaN

问题

使用tf.keras.applications.ResNet101V2作为backbone的多输出多分类模型训练正常，推理时输出nan。使用tf混合精度进行训练，期间有出现过loss nan，降低学习率后解决。在差不多收敛后进行predict测试，发现模型输出nan。

调试

检查模型权重

输出模型weight，搜索nan，发现某些层的权重存在nan的情况，具体来说大部分集中在最后一个BatchNormalization层。

def check_and_initialize_nan_weights(model):
    for layer in model.layers:
        for weight in layer.weights:
            try:
                nan_indices = tf.debugging.check_numerics(weight, "NaN detected in weight")
            except Exception as e:
                # assert "Checking b : Tensor had NaN values" in e.message
                print("NaN detected in weight:", weight.name)
                # layer.kernel_initializer(shape=np.asarray(layer.kernel.shape)), \
                #                    layer.bias_initializer(shape=np.asarray(layer.bias.shape))])
                # 对 NaN 权重进行初始化
                nan_mask = tf.math.is_nan(weight)
                weight.assign(tf.where(nan_mask, tf.ones_like(weight), weight))

检查中间层输出

有时候不仅仅是权重问题，一些计算逻辑也可能导致nan，于是查看了模型中间层的输出，找到到底是哪一层计算逻辑后出现了nan，结果是最后一个resnet block中的命名为conv5_block3_2_bn的BatchNormalization层的moving_variance过大，以及命名为post_bn的全局池化层之前的最后一个BatchNormalization层的moving_mean和moving_variance都出现了nan。

res = model.get_layer('resnet101v2')
bn = res.get_layer('conv5_block3_1_bn')
post_bn = res.get_layer('post_bn')
bn_2 = res.get_layer('conv5_block3_2_bn')
bn_2.set_weights([bn_2.weights[0],bn_2.weights[1],bn_2.moving_mean_initializer(shape=np.asarray(bn_2.moving_mean.shape)),bn_2.moving_variance_initializer(shape=np.asarray(bn_2.moving_variance.shape))])
post_bn.set_weights([post_bn.weights[0],post_bn.weights[1],post_bn.moving_mean_initializer(shape=np.asarray(post_bn.moving_mean.shape)),post_bn.moving_variance_initializer(shape=np.asarray(post_bn.moving_variance.shape))])
input = tf.keras.Input(shape=(224, 224, 3), name="input")
sub = keras.Model(inputs=res.input, outputs=res.get_layer('conv5_block3_out').output)
output = sub(input)
submodel = keras.Model(inputs=input, outputs=output)
submodel.summary()
print(submodel(image))

分析

BatchNormalization

此次模型输出nan是由于BatchNormalization层参数异常造成的，至于为什么训练时候正常是由于BatchNormalization层在训练与推理时不同的行为造成的。

BatchNormalization应用一种变换，该变换将平均输出（均值）保持在接近0的水平，并将输出标准偏差保持在接近1的水平。批量归一化在训练和推理期间的工作方式不同。

在训练期间（即使用fit()或使用参数调用层/模型时training=True），层使用当前批次输入的平均值和标准差标准化其输出。也就是说，对于每个被归一化的通道，该层返回 gamma * (batch - mean(batch)) / sqrt(var(batch) + epsilon) + beta，其中：

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

在推理过程中（即，当使用evaluate()或predict()或使用参数调用层/模型时training=False（这是默认值），该层使用在训练期间看到的批次的平均值和标准差的移动平均值来规范化其输出。即说，它回来了 gamma * (batch - self.moving_mean) / sqrt(self.moving_var+epsilon) + beta。

self.moving_mean和self.moving_var是不可训练的变量，每次在训练模式下调用层时都会更新，如下所示：