TensorFlow2.5.0使用tfp.layers.Convolution3DFlipout报Graph张量错误是什么原因

• 核心原因是KL散度的计算时机错误：你当前的代码在返回损失函数对象之前，就提前在外层variational_free_energy_loss函数中执行了sum(model.losses) / scale_factor的计算，这个计算发生在Keras训练时的tf.function图构建上下文之外，得到的kl是独立于训练步骤动态计算图的静态Graph张量。
• TFP的Convolution3DFlipout这类变分层的KL散度损失项，会在层前向传播时自动注册到model.losses列表中，你提前在损失函数逻辑之外求和得到的kl张量，和后续训练时每步动态生成的损失项不属于同一个计算图，就会触发图张量泄露的报错。
• 你使用的TensorFlow 2.5.0 + TensorFlow Probability 0.13.0版本对损失函数闭包的张量捕获规则比较严格，闭包中引用的外部静态图张量无法被训练时的动态计算图正确识别，也会加剧这个问题的出现。

把KL散度的计算移到内层损失函数中，保证每步训练时KL都在当前计算图上下文内生成：

from tensorflow.keras.losses import binary_crossentropy

def variational_free_energy_loss(model, scale_factor = tf.constant(1.)):
    def loss(y_true, y_pred):
        kl = sum(model.losses) / scale_factor
        bce = binary_crossentropy(y_true, y_pred)
        return bce + kl 
    return loss

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