TensorFlow 2.2 Eager模式下如何获取梯度？model.total_loss已弃用

你遇到的问题是TensorFlow 2.2在Eager模式下移除了model.total_loss的直接访问，下面我会给出完全兼容learning_phase标志和sample_weight的替代方案，同时匹配原代码的核心功能。

核心思路 #

在TF2.x的Eager模式下，我们可以通过tf.GradientTape追踪梯度计算，同时复用模型编译好的损失逻辑（保证和model.compile()的设置一致），手动处理样本权重和训练/推断模式的切换。如果需要兼容Graph模式，也可以用K.function包装实现。

完整实现代码 #

方案1：Eager模式优先（推荐） #

import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model
from tensorflow.keras import backend as K

# 1. 构建并编译模型（和原代码一致）
ipt = Input((16,))
out = Dense(16)(ipt)
model = Model(ipt, out)
model.compile('adam', 'mse')

# 2. 准备测试数据（包含sample_weight示例）
x = y = np.random.randn(32, 16)
sample_weight = np.random.rand(32,)  # 随机生成样本权重

# 3. 定义梯度获取函数，支持learning_phase和sample_weight
def get_model_gradients(model, x, y, sample_weight=None, training=True):
    # 设置learning_phase，控制Dropout/BatchNorm等层的训练行为
    K.set_learning_phase(training)
    
    with tf.GradientTape() as tape:
        # 前向传播，training参数显式控制训练模式
        y_pred = model(x, training=training)
        # 复用模型编译好的损失计算逻辑
        loss = model.compiled_loss(y, y_pred)
        # 处理样本权重，和Keras内部逻辑对齐
        if sample_weight is not None:
            loss = tf.reduce_mean(loss * sample_weight)
    
    # 计算并返回梯度
    gradients = tape.gradient(loss, model.trainable_weights)
    return gradients

# 4. 测试获取梯度
grad_tensors = get_model_gradients(model, x, y, sample_weight=sample_weight)

方案2：兼容Graph模式（类似原代码的K.function实现） #

如果你需要保留原代码中用K.function构建计算图的方式，也可以这样实现：

def get_grad_function(model):
    # 定义输入：模型输入、标签、样本权重、learning_phase标志
    inputs = [model.input, model.targets[0], model.sample_weights[0], K.learning_phase()]
    
    def compute_loss(x, y, sw, training):
        y_pred = model(x, training=training)
        loss = model.compiled_loss(y, y_pred)
        if sw is not None:
            loss = tf.reduce_mean(loss * sw)
        return loss
    
    # 计算梯度
    grads = K.gradients(compute_loss(*inputs), model.trainable_weights)
    return K.function(inputs, grads)

# 使用示例：1表示training模式，0表示inference模式
grad_fn = get_grad_function(model)
gradients_from_fn = grad_fn([x, y, sample_weight, 1])

关键细节说明 #

• 复用compiled_loss：直接调用model.compiled_loss能保证损失计算逻辑和你model.compile()时指定的完全一致，避免手动实现损失带来的偏差。
• 样本权重处理：通过tf.reduce_mean(loss * sample_weight)缩放损失，和Keras内部处理样本权重的逻辑完全对齐，确保梯度计算准确。
• learning_phase控制：两种方案都支持通过参数切换训练/推断模式，确保Dropout、BatchNormalization等层的行为符合预期。
• Eager模式优势：方案1的tf.GradientTape是TF2.x的标准梯度获取方式，更直观且兼容Eager模式的动态计算特性。

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