Keras含外部数据的自定义损失函数无梯度报错问题求助

你的代码出现梯度缺失的问题，核心原因有三个：

1. 不可导操作切断梯度路径 #

math_ops.argmax(y_pred, axis=0)是不可微分的操作——它返回离散的索引值，TensorFlow无法为该操作计算梯度，导致后续所有计算和模型参数（LSTM、Dense的权重）之间的梯度连接完全断开。

2. 原地修改外部变量破坏计算图 #

classes_slice.assign(_class)直接原地修改外部定义的K.variable，这种操作不属于TensorFlow计算图的一部分，无法被自动微分机制追踪。而且训练过程中批次之间会互相覆盖这个变量的状态，不仅导致梯度失效，还会引入不可控的训练波动。

3. 损失返回逻辑的梯度断裂 #

虽然tf.fill((A,), H)的形状和示例一致，但H的计算依赖前两个无梯度的操作，最终导致H和模型参数之间没有任何梯度路径，优化器自然拿不到参数的梯度。

修复方案 #

需要重构损失函数，用可微分操作替代不可导步骤，同时避免修改外部状态：

关键调整点： #

• 用可微分的加权逻辑替代argmax：直接用y_pred的概率值作为权重更新分布，而非取离散索引。
• 放弃原地修改外部变量：所有运算基于张量正向传播，保持损失函数的无状态性。
• 损失返回标量或与批次样本数匹配的张量：Keras默认会对批次内损失取平均，无需手动填充重复值。

修改后的代码示例： #

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
import numpy as np

A = 35  # 假设你的A取值为35
data = np.random.random((100, A, 1))  # 示例输入数据

# 初始化分布为常量张量，避免原地修改
initial_distributions = np.full((A, A), 1.0 / A, dtype=np.float32)

def my_loss():
    def loss(y_true, y_pred):
        # 用y_pred的概率分布作为权重（可微分），替代argmax的离散索引
        y_true_flat = tf.reshape(y_true, (-1, A))
        y_pred_soft = tf.nn.softmax(y_pred, axis=1)  # 确保输出为合法概率分布
        
        # 用张量运算实现分布更新逻辑，无原地修改
        updated_dist = initial_distributions + tf.matmul(tf.transpose(y_pred_soft), y_true_flat)
        updated_dist = updated_dist / tf.reduce_sum(updated_dist, axis=1, keepdims=True)
        
        # 计算平均熵，加小值避免log(0)报错
        log_p = tf.math.log(updated_dist + 1e-10)
        H = -tf.reduce_sum(updated_dist * log_p, axis=1)
        mean_H = tf.reduce_mean(H)
        
        # 返回标量损失，Keras自动处理批次平均
        return mean_H
    return loss

model = Sequential()
model.add(LSTM(units=A*10, input_shape=(A, 1)))
model.add(Dense(units=A, activation='softmax'))
model.compile(loss=my_loss(), optimizer='adam')
history = model.fit(
    data, 
    data,
    epochs=200,
    verbose=True, 
    batch_size=A)

额外说明： #

如果业务逻辑必须依赖离散类别选择（类似argmax的行为），可以使用Gumbel-Softmax技巧近似离散选择的可微分版本，避免梯度消失。

内容的提问来源于stack exchange，提问作者Дмитрий Шихов