Adam优化器因tf.assign函数引发梯度为None的错误及解决方法

问题场景 #

当使用input = [self.u, self.v]配合当前loss_total计算逻辑时，触发梯度为None的错误，报错信息如下：

ValueError: No gradients provided for any variable: (['Variable_541:0', 'Variable_542:0'],). Provided `grads_and_vars` is ((None, <tf.Variable 'Variable_541:0' shape=(5, 5) dtype=float32>), (None, <tf.Variable 'Variable_542:0' shape=(5, 5) dtype=float32>)).

报错发生在optimizer_utils.filter_empty_gradients函数中。

错误原因 #

核心问题是loss_estimation方法里的原地赋值操作assign切断了梯度流：

self.U_padded[1:-1,1:-1].assign(self.u)
self.V_padded[1:-1,1:-1].assign(self.v)

TensorFlow的tf.GradientTape无法追踪原地修改变量的操作，导致loss_total（基于U_padded和V_padded计算）与原始变量u、v之间没有可微分的计算链路，最终无法计算出梯度。

解决方案 #

替换原地assign操作，改用可微分的张量填充逻辑，让U_padded、V_padded成为u、v的派生张量，从而保留梯度传递链路。

修改后的代码 #

1. 重写loss_estimation方法

移除assign，用tf.pad实现填充逻辑：

@tf.function
def loss_estimation(self):
    # 用tf.pad生成带padding的张量，替代原地assign
    self.U_padded = tf.pad(self.u, [[1, 1], [1, 1]], mode='CONSTANT')
    self.V_padded = tf.pad(self.v, [[1, 1], [1, 1]], mode='CONSTANT')
    self.loss_total = mean_squared_error(self.U_padded + self.V_padded, 0)
    return self.loss_total

2. 保留原train方法不变

@tf.function
def train(self, optimizer):
    with tf.GradientTape() as tape:
        input = [self.u, self.v]
        self.loss = self.loss_estimation()
    self.gradients = tape.gradient(self.loss, input)
    optimizer.apply_gradients(zip(self.gradients, input))

额外说明 #

如果必须保留U_padded、V_padded为类变量，无需在__init__中初始化它们，直接在loss_estimation中动态生成即可，这样既满足计算需求，又能保证梯度正常传递。

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