TensorFlow自定义角距离损失训练数个epoch后Loss为NaN求助

• 梯度爆炸问题：tf.math.acos的导数为 -1 / sqrt(1 - x²)，当点积v_dot的取值非常接近±1时，分母趋近于0，梯度会变为无穷大，导致参数更新溢出，最终产生NaN值。即使你做了clip到±1的操作，刚好在边界处的梯度依然是发散的。
• 输出范围无约束问题：你仅对theta取绝对值、对phi做模运算，但没有限制输出的取值范围，当模型预测的theta、phi出现极大/极小的异常值时，三角函数计算后可能引入数值不稳定，同时模运算在输入为极端值时也可能出现精度问题。
• 非核心但不规范的维度处理：angle_dist已经是每个样本对应的标量损失，K.mean(K.square(angle_dist), axis=-1)的axis=-1没有实际作用，虽然不会直接导致NaN，但存在维度误用的风险。

1. 替换角距离计算方式，使用数值更稳定的arctan2实现，避免acos的梯度爆炸问题：角距离可以通过叉积的模和点积计算：Δ = arctan2(||v_true × v_pred||, v_true·v_pred)，该公式的梯度在所有合法输入下都保持稳定。
2. 约束输出层的取值范围：将输出层的两个神经元分别绑定激活函数，不需要再用模运算和取绝对值：
   • theta输出用sigmoid激活后乘以π，直接映射到0~π的合法区间
   • phi输出用sigmoid激活后乘以2π，直接映射到0~2π的合法区间
3. 点积clip时保留微小余量，避免刚好触碰到±1的边界：将clip范围调整为(-0.999999, 0.999999)

import tensorflow as tf
import numpy as np
import tensorflow.keras.backend as K

def square_angular_distance(y_true, y_pred):
    # 输出层已经用激活约束范围后，不需要abs和模运算
    the_pred = y_pred[:, 0]
    phi_pred = y_pred[:, 1]

    the_true = y_true[:, 0]
    phi_true = y_true[:, 1]

    cos_phi_pred = K.cos(phi_pred)
    cos_phi_true = K.cos(phi_true)
    sin_phi_pred = K.sin(phi_pred)
    sin_phi_true = K.sin(phi_true)

    cos_the_pred = K.cos(the_pred)
    cos_the_true = K.cos(the_true)
    sin_the_pred = K.sin(the_pred)
    sin_the_true = K.sin(the_true)

    v_true = K.stack((sin_the_true*cos_phi_true, sin_the_true*sin_phi_true, cos_the_true), axis=1)
    v_pred = K.stack((sin_the_pred*cos_phi_pred, sin_the_pred*sin_phi_pred, cos_the_pred), axis=1)

    v_dot = K.clip(K.batch_dot(v_true, v_pred), -0.999999, 0.999999)
    # 用叉积+arctan2计算稳定的角距离
    v_cross = tf.linalg.cross(v_true, v_pred)
    cross_norm = K.sqrt(K.sum(K.square(v_cross), axis=-1))
    angle_dist = tf.math.atan2(cross_norm, v_dot) * 180. / np.pi

    return K.mean(K.square(angle_dist))

# 输出层定义示例，直接约束角度范围，不需要后续在损失里做abs和模运算
output = tf.keras.layers.Dense(2, activation=lambda x: tf.stack([
    tf.sigmoid(x[:, 0]) * np.pi,  # theta约束到0~π
    tf.sigmoid(x[:, 1]) * 2 * np.pi  # phi约束到0~2π
], axis=-1))(prev_layer)

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