TensorFlow自定义损失函数复制y_pred后失效问题排查

该问题不是tf.identity的张量复制逻辑错误。tf.identity本身的功能就是返回与输入张量数值完全一致的新张量，数值计算层面没有问题。故障根源是TensorFlow/Keras计算图追踪与自动求导的特殊行为：

• 直接使用训练循环传入的y_pred计算损失时，Keras可以正确识别该张量是模型前向传播的输出节点，自动建立从损失值到模型所有可训练权重的完整反向传播路径，梯度可以正常回传，因此模型能正常学习。
• 使用裸tf.identity(y_pred)生成新张量计算损失时，在TensorFlow 2.x早期版本（2.0_{2.4区间）中，新生成的张量节点会脱离Keras对模型输出节点的追踪链路，导致反向传播阶段梯度无法传递回模型权重。你观察到的训练停滞、AUC恒为0.5，本质就是模型权重从始至终没有更新，初始状态下模型输出经sigmoid激活后接近0.5，对应的二分类交叉熵损失值正好稳定在1.25}1.26区间，和日志记录的数值完全吻合。

不要直接使用无参数的裸tf.identity处理损失函数入参y_pred，可选择以下不会中断梯度链路的复制方式：

• 给tf.identity显式指定名称，避免被图优化逻辑错误截断梯度路径：

  y_p = tf.identity(y_pred, name="prediction_copy")
  

• 使用不改变张量数值的逐元素运算生成副本，这类运算的梯度传播逻辑完全透明，不会被图优化误判：

  # 乘1、加0均为数值恒等运算，返回结果就是原张量的独立副本
  y_p = y_pred * 1.0
  y_p = y_pred + 0.0
  

• 使用tf.ensure_shape在做形状校验的同时完成复制，同样不会中断梯度：

  y_p = tf.ensure_shape(y_pred, y_pred.shape)
  

如果需要生成完全独立、后续要做掩码/截断等修改操作的副本，最稳妥的方式是把复制逻辑放在模型前向传播过程中（比如在模型输出层后加一个Lambda层完成复制），让复制生成的张量本身就是模型计算图的固定节点，从根源上避免追踪断链问题。

你可以在训练时打印模型权重的梯度值做验证：异常版本下所有权重的梯度均为0，参数完全不更新；正常版本梯度非零，参数可以正常迭代优化。

内容的提问来源于stack exchange，提问作者Athreya H P