使用自定义损失函数训练TensorFlow模型时为何无法计算梯度？

问题背景 #

损失函数可正常运行，但无法计算梯度。batter_cdfs是一个2150×10的浮点数数组，数值范围在0到1之间。该损失函数的目标是返回模型生成的联合分布中单个观测样本的对数概率。

原始代码 #

import tensorflow as tf
import numpy as np
import pandas as pd

def custom_loss(y_true, y_pred):
    y_true_expanded = tf.expand_dims(y_true, axis=1)
    y_true_expanded = tf.tile(y_true_expanded, [1, tf.shape(y_pred)[1], 1])
    
    mask = tf.less(y_pred, y_true_expanded)
    
    tf.print("mask:", mask)
    
    row_wise_comparison = tf.reduce_mean(tf.cast(mask, tf.float32), axis=2)
    
    result = tf.reduce_mean(row_wise_comparison, axis=1)
    
    log_result = tf.math.log(result + 1e-9)  # 加小值避免log(0)
    
    return -tf.reduce_mean(log_result)

cdf_model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='LeakyReLU'),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dropout(0.1),
    
    tf.keras.layers.Dense(4500, activation='sigmoid'),
    tf.keras.layers.Reshape((500,9))
])

cdf_model.compile(optimizer=tf.optimizers.Nadam(), loss=custom_loss)
input_shape = (100,)
# 假设batter_cdfs是pd.DataFrame
y_train = np.array(batter_cdfs.iloc[:,0:9])
x_train = np.array(batter_cdfs.iloc[:,9:])
cdf_model.fit(x_train, y_train, epochs=15, validation_split=0.2, verbose=True, shuffle=True, batch_size=30)

问题原因 #

核心问题出在tf.less(y_pred, y_true_expanded)这一步：

• 该操作生成布尔张量，转成float32后等价于阶跃函数，在y_pred == y_true_expanded处导数不存在，其余位置导数为0。
• 反向传播时，梯度无法通过这个硬阈值操作传递到前面的模型参数，导致模型无法更新。

修改方案 #

用平滑近似函数替代硬比较操作，让梯度能正常传播。这里选择sigmoid函数来近似阶跃行为，调整参数控制平滑程度：

def custom_loss(y_true, y_pred):
    y_true_expanded = tf.expand_dims(y_true, axis=1)
    y_true_expanded = tf.tile(y_true_expanded, [1, tf.shape(y_pred)[1], 1])
    
    # 用sigmoid近似硬比较，100是平滑系数，数值越大越接近阶跃
    mask = tf.sigmoid(-100 * (y_pred - y_true_expanded))
    
    row_wise_comparison = tf.reduce_mean(mask, axis=2)
    
    # 限制result下限，避免log(0)
    result = tf.clip_by_value(tf.reduce_mean(row_wise_comparison, axis=1), 1e-9, 1.0)
    
    log_result = tf.math.log(result)
    
    return -tf.reduce_mean(log_result)

额外说明 #

• 平滑系数（示例中的100）可根据需求调整：数值越大，mask越接近原始的硬比较，但梯度在阈值附近的变化越剧烈；数值越小，平滑度越高，但近似精度会下降。
• 使用tf.clip_by_value替代直接加1e-9，能更稳定地避免log(0)问题。
• 训练时可去掉tf.print语句，避免冗余输出影响效率。

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