TensorFlow 2.5推理时如何同时启用Dropout并禁用BatchNormalization

现有写法效果 #

你的当前写法无法在推理阶段保持Dropout启用：TensorFlow中调用模型时传入的training参数优先级高于层初始化时单独设置的training固定值，你传入training=False时会覆盖Dropout层的training=True设置，直接导致推理时Dropout被关闭。

可用实现方案 #

下面两种方案都可以实现「训练阶段BN和Dropout正常工作，推理阶段Dropout保持启用、BN切换为推理模式」的需求：

方案1：自定义始终启用的Dropout层（推荐，训练推理通用） #

直接自定义Dropout子类，强制忽略全局training参数始终启用dropout，替换原有模型中的原生Dropout层即可：

import tensorflow as tf

class AlwaysOnDropout(tf.keras.layers.Dropout):
    def call(self, inputs, training=None):
        # 强制启用dropout，忽略外部传入的training参数
        return super().call(inputs, training=True)

• 训练阶段：正常调用模型训练即可，BN会处于训练模式更新滑动均值和方差，Dropout也正常生效
• 推理阶段：直接调用model(X_test, training=False)或者model.predict(X_test)即可，BN会切换为推理模式使用训练阶段统计的滑动值，自定义的Dropout层会忽略全局的training=False保持启用。

方案2：基于已训练模型生成推理专用副本（无需重训） #

如果你已经用原生Dropout完成训练，不想修改训练代码或重训，可以直接遍历训练好的模型层，单独控制两类层的推理行为，生成专用的推理模型：

def build_inference_model(trained_model):
    inputs = trained_model.input
    x = inputs
    for layer in trained_model.layers:
        if isinstance(layer, tf.keras.layers.Dropout):
            # Dropout强制开启
            x = layer(x, training=True)
        elif isinstance(layer, tf.keras.layers.BatchNormalization):
            # BN强制关闭训练模式
            x = layer(x, training=False)
        else:
            x = layer(x)
    return tf.keras.Model(inputs=inputs, outputs=x)

# 生成推理模型后直接调用即可
inference_model = build_inference_model(trained_model)
prediction = inference_model(X_test)

内容的提问来源于stack exchange，提问作者axel