Python闭包向Keras自定义损失函数传递标量值失效问题

出现该问题是两个机制共同导致的：

1. 打印时机错误：你在损失函数里用的是Python原生print，Keras/TensorFlow默认采用静态图执行模式，原生Python语句只会在计算图构建阶段执行1次，不会在后续每轮训练计算损失时重复运行。你看到的输出0是图构建阶段捕获的初始值，不是训练过程中损失函数实际拿到的实时值。
2. 值被静态图固化：Python原生的整数、浮点数属于不可变类型，当你把普通标量传入闭包时，TensorFlow在图构建阶段会直接把这个初始值固化为计算图里的常量，后续你在外部修改add_loss2变量的取值，完全不会影响已经构建完成的计算图里的常量值，所以损失计算时实际用的永远是图构建时传入的初始值（也就是你看到的0）。

用TensorFlow可追踪的tf.Variable承载动态变化的标量，同时替换原生打印语句为图执行阶段可运行的tf.print，具体实现如下：

1. 首先把需要动态调整的add_loss2定义为非训练的tf.Variable，注意数据类型要和模型输出、标签的浮点类型保持一致：

import tensorflow as tf
from tensorflow.keras import losses

# 初始化可变变量，初始值可按需设置，这里以0.0为例
add_loss2 = tf.Variable(initial_value=0.0, dtype=tf.float32, trainable=False)

2. 修改闭包实现，把原生print替换为tf.print，直接接收传入的Variable对象：

@staticmethod
def sdec_loss(add_loss_var):
    def loss_function(y_true, y_pred):
        tf.print("add_loss2= ", add_loss_var) # 每次计算损失都会打印实时值
        return losses.kullback_leibler_divergence(y_true, y_pred) - add_loss_var
    return loss_function

3. 绑定损失时直接传入定义好的Variable对象即可：

self.DEC.loss = SDEC.sdec_loss(add_loss2)

4. 后续训练过程中需要调整add_loss2的取值时，调用Variable的assign方法更新，损失函数会自动读取最新值，不需要重新绑定损失：

# 示例：将add_loss2更新为0.3
add_loss2.assign(0.3)

• 不要在Keras/TF的损失函数、层的call方法里写原生Python的打印、条件判断、循环这类语句，这类语句只会在图构建时执行一次，不会在训练运行时生效，需要对应使用tf.print、tf.cond、tf.while_loop等TF提供的可图化算子。
• 如果你的add_loss2是每个batch都变化的动态值，除了用Variable更新的方案，也可以把它作为模型的额外输入传入，在损失计算时调用，但这种方案需要修改模型的输入结构，适合每个batch取值都不同的场景。

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