Tensorflow自定义tf.function训练报No gradients provided for any variable问题

报错原因 #

这个错误本质是计算图中损失值和模型可训练变量之间没有可微分的路径，你的代码存在两个容易导致梯度截断的问题：

1. 手动传入类对象model到tf.function修饰的独立函数时，GradientTape默认不会自动监视外部传入模型的可训练变量，导致梯度追踪失败
2. rewards为整数类型、actions类型和tf.range不匹配的隐性类型转换，以及损失没有做降维处理，会提升计算图构建异常的概率

修复方案 #

方案1：调整为类方法封装（推荐） #

不要把模型、优化器作为tf.function的参数传入，建议将训练逻辑封装为类方法，在类初始化阶段完成tf.function的绑定，避免Python对象参数带来的追踪异常：

class DQNTrainer:
    def __init__(self, model, optimizer):
        self.model = model
        self.optimizer = optimizer
        # 类初始化阶段绑定tf.function，避免self参数带来的重复追踪问题
        self.train_step = tf.function(self._train_step_impl)
    
    def _train_step_impl(self, states, actions, rewards):
        with tf.GradientTape() as tape:
            # 强制监听模型可训练变量，避免梯度追踪遗漏
            tape.watch(self.model.trainable_variables)
            predictions = self.model(states)
            
            # 统一输入张量类型，避免隐性类型转换
            actions = tf.cast(actions, tf.int64)
            rewards = tf.cast(rewards, tf.float32)
            indices = tf.stack([tf.range(predictions.shape[0], dtype=tf.int64), actions], axis=1)
            action_rewards = tf.scatter_nd(indices, rewards, predictions.shape)
            
            # 损失降维为标量，符合优化器默认输入要求
            loss = tf.reduce_mean(tf.keras.losses.MSE(predictions, action_rewards))
        
        grads = tape.gradient(loss, self.model.trainable_variables)
        # 可选增加梯度裁剪，避免训练不稳定
        grads, _ = tf.clip_by_global_norm(grads, 5.0)
        self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
        return loss

调用时提前将输入转换为张量类型，避免Python列表/NumPy数组带来的计算图重追踪问题：

# 调用示例
states_tensor = tf.convert_to_tensor(states_epoch, dtype=tf.float32)
actions_tensor = tf.convert_to_tensor(actions_epoch, dtype=tf.int64)
rewards_tensor = tf.convert_to_tensor(rewards_epoch, dtype=tf.float32)
loss = trainer.train_step(states_tensor, actions_tensor, rewards_tensor)

方案2：保留独立函数封装 #

如果确实需要将训练逻辑封装为独立函数，需要在GradientTape中手动添加对模型可训练变量的监听，同时对所有输入做类型统一：

@tf.function
def train_network(model, optimizer, states, actions, rewards):
    with tf.GradientTape() as tape:
        # 强制监听模型可训练变量
        tape.watch(model.trainable_variables)
        predictions = model(states)
        actions = tf.cast(actions, tf.int64)
        rewards = tf.cast(rewards, tf.float32)
        indices = tf.stack([tf.range(predictions.shape[0], dtype=tf.int64), actions], axis=1)
        action_rewards = tf.scatter_nd(indices, rewards, predictions.shape)
        loss = tf.reduce_mean(tf.keras.losses.MSE(predictions, action_rewards))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

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