tf.keras的model.fit如何让批次内每个样本被推理n次以适配MC dropout损失

核心思路 #

通过输入样本重复推理+输出维度重排的方式，得到每个样本的n次MC dropout预测结果，再调整损失函数的统计维度，即可得到每个样本专属的均值、方差。

完整修改代码 #

import numpy as np
import tensorflow as tf

# 定义超参数
n = 5 # 自定义MC采样次数，可按需调整
X = np.random.rand(5000,24)
y = np.random.rand(5000,10)

def MC_Loss(y_true,y_pred):
    # 扩展y_true维度适配广播规则
    y_true = tf.expand_dims(y_true, axis=1)
    # 沿采样维度统计每个样本的均值、方差
    mu = tf.math.reduce_mean(y_pred,axis=1) # 形状 (None, 10)
    var = tf.math.reduce_variance(y_pred,axis=1) + 1e-6 # 加小常数避免除以0
    error = tf.square(y_true - mu)
    return tf.math.divide(error,var)/2 + tf.math.log(var)/2 + tf.math.log(2*np.pi)/2

input_layer = tf.keras.layers.Input(shape=(X.shape[1],))
# 每个样本重复n次，保证每次推理生成不同的dropout掩码
repeat_input = tf.keras.layers.Lambda(lambda x: tf.repeat(x, repeats=n, axis=0))(input_layer)
hidden_layer = tf.keras.layers.Dense(units=100,activation='elu')(repeat_input)
do_layer = tf.keras.layers.Dropout(rate=0.20)(hidden_layer,training=True)
output_layer = tf.keras.layers.Dense(units=10,activation='sigmoid')(do_layer)
# 重排输出维度为 (批次大小, 采样次数n, 输出维度10)
reshape_output = tf.keras.layers.Lambda(lambda x: tf.reshape(x, (-1, n, 10)))(output_layer)

model = tf.keras.models.Model(input_layer, reshape_output)
model.compile(loss=MC_Loss,optimizer='Adam')

# 验证输出形状是否符合预期
test_batch = np.random.rand(128, 24)
pred = model.predict(test_batch, verbose=0)
print(f"输出形状：{pred.shape}，符合 (None, {n}, 10) 的预期")

model.fit(X,y,epochs=100,batch_size=128,shuffle=True)

关键修改说明 #

• 采样次数配置：参数n可根据精度和速度需求调整，采样次数越多不确定性估计越稳定，但显存占用和训练耗时也会同步提升
• 输入重复逻辑：通过tf.repeat将每个样本在批次维度复制n次，每次推理时dropout生成独立随机掩码，得到n个互不干扰的MC预测结果
• 维度重排：推理完成后将输出从(批次大小*n, 10)重排为(批次大小, n, 10)，匹配损失函数的输入要求
• 损失维度调整：将均值、方差的统计维度从原代码的axis=0改为axis=1，即可得到每个样本独立的统计量，输出形状为(None, 10)，完全符合需求
• 数值稳定性优化：方差计算后添加1e-6的小常数，避免出现除以0的数值报错

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