TensorFlow V1转V2中Adam.minimize与GradientTape报错问题

问题背景 #

在将TensorFlow V1代码迁移到V2时，使用tf.optimizers.Adam.minimize遇到梯度磁带相关错误，尝试添加磁带参数后仍无法解决，怀疑与类中已存在的GradientTape冲突。

复现代码 #

import tensorflow as tf
import numpy as np

args = {'layers': {'j': [4, 6, 8], 'q': [8, 2, 1], 'k': 7, 'pooling_type': 'average', 'decimation': 4, 'pooling': 1024, 'learn_scales': False, 'learn_knots': False, 'learn_filters': True, 'hilbert': True}, 'gmm': {'gmm_type': 'natural', 'trainable': False}, 'gmm_init': {'n_components': 10, 'max_iter': 1000, 'covariance_type': 'full', 'warm_start': True}}

class temp:
    def __init__(self, x, j=None, q=None, k=None, pooling_type='average',
                 decimation=2, pooling=2, index=0, **filters_kw):
        # Filter bank properties
        with tf.GradientTape() as tape:
          tape.watch(x)
          x_conv = tf.random.uniform((12,64,4201), dtype=tf.float32)
          # 完整代码中的其他逻辑

        inverse = tape.gradient(x_conv, x, output_gradients=x_conv) # 完整代码中此处不会返回None
        inverse = tf.random.uniform((4,3,16800), dtype=tf.float32) # 临时赋值让代码运行
        reconstruction_loss = tf.nn.l2_loss(inverse - tf.stop_gradient(x))

data = tf.random.uniform((4,3,16800), dtype=tf.float32)
layers = [temp(data, index=0, **args['layers'])]

# 原始loss定义方式：gmm返回浮点型loss值
# loss, cat = gmm(mu, cov, tau, sx_proj, **args['gmm'])
loss = tf.convert_to_tensor(1.123)

optimizer = tf.optimizers.Adam(0.001)
sx_proj = tf.random.uniform((64,5), dtype=tf.float32)
tau = np.random.rand(10,)

minimizer = optimizer.minimize(loss, var_list=[tau, sx_proj], tape=tf.GradientTape(persistent=True)) 

错误信息 #

初始错误（未添加tape参数） #

ValueError: `tape` is required when a `Tensor` loss is passed. Received: loss=1.1230000257492065, tape=None.

添加磁带后新错误 #

RuntimeError                              Traceback (most recent call last)
<ipython-input-15-82d22d49e786> in <module>
     27 optimizer = tf.optimizers.Adam(0.001)
     28 
---> 29 minimizer = optimizer.minimize(loss, var_list=[tau, sx_proj], tape=tf.GradientTape(persistent=True))
     30 
     31 # SOS , tape=tf.GradientTape(persistent=True)

2 frames
/usr/local/lib/python3.8/dist-packages/keras/optimizers/optimizer_experimental/optimizer.py in minimize(self, loss, var_list, tape)
    524           None
    525         """
--> 526         grads_and_vars = self.compute_gradients(loss, var_list, tape)
    527         self.apply_gradients(grads_and_vars)
    528 

/usr/local/lib/python3.8/dist-packages/keras/optimizers/optimizer_experimental/optimizer.py in compute_gradients(self, loss, var_list, tape)
    257                     var_list = var_list()
    258 
--> 259         grads = tape.gradient(loss, var_list)
    260         return list(zip(grads, var_list))
    261 

/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/backprop.py in gradient(self, target, sources, output_gradients, unconnected_gradients)
   1050     """
   1051     if self._tape is None:
-> 1052       raise RuntimeError("A non-persistent GradientTape can only be used to "
   1053                          "compute one set of gradients (or jacobians)")
   1054     if self._recording:

RuntimeError: A non-persistent GradientTape can only be used to compute one set of gradients (or jacobians)

相关参考代码（原始GMM函数） #

import tensorflow_probability as tfp

tfd = tfp.distributions
MVNFC = tfd.MultivariateNormalFullCovariance

def gmm(mu, cov, tau, sx_proj, n_clusters=None, gmm_type='natural',
        trainable=True, cov_diag=None):
    """Gaussian mixture clustering."""
    if trainable is True:
        log_tau = tf.log(tf.nn.softmax(tau))
        eps = tf.eye(2)
        d = cov_diag
        gm = [MVNFC(mu[c], tf.nn.elu(d[c]) * np.eye(2) +
                    tf.matmul(cov[c], tf.transpose(cov[c])) + eps)
              for c in range(n_clusters)]
    log_p = [gm[c].log_prob(sx_proj) for c in range(n_clusters)]
    cat = tf.stack([log_tau[c] + log_p[c] for c in range(n_clusters)], 1)

    # 选择loss类型
    if gmm_type == 'natural':
        q = tf.math.reduce_logsumexp(cat, axis=1)
        loss = - tf.reduce_mean(q)

    return loss, cat

解决方案 #

核心问题分析 #

1. 手动传入tf.GradientTape(persistent=True)时，新创建的磁带并未记录loss与待优化变量tau、sx_proj之间的计算过程，且非持久化磁带只能使用一次。
2. 类temp中的磁带仅记录了自身内部的计算，和外部优化的变量无关。

正确实现方式 #

需要在计算loss的整个过程外包裹梯度磁带，确保磁带记录了从待优化变量到loss的完整计算链路：

import tensorflow as tf
import numpy as np

# 保留原args、temp类、gmm函数定义不变...

data = tf.random.uniform((4,3,16800), dtype=tf.float32)
layers = [temp(data, index=0, **args['layers'])]

optimizer = tf.optimizers.Adam(0.001)
# 将tau转为TensorFlow变量，确保可被优化
tau = tf.Variable(np.random.rand(10,), dtype=tf.float32)
sx_proj = tf.Variable(tf.random.uniform((64,5), dtype=tf.float32))

# 关键：在计算loss的过程外包裹GradientTape
with tf.GradientTape() as tape:
    # 此处替换为真实的loss计算逻辑，确保tau和sx_proj参与计算
    # loss, cat = gmm(mu, cov, tau, sx_proj, **args['gmm'])
    loss = tf.convert_to_tensor(1.123)  # 仅示例，需替换为真实计算

# 计算梯度并应用
grads = tape.gradient(loss, [tau, sx_proj])
optimizer.apply_gradients(zip(grads, [tau, sx_proj]))

关键注意点 #

• 待优化变量必须是tf.Variable类型，而非numpy数组（原代码中tau是numpy数组，无法被梯度磁带追踪）。
• 梯度磁带必须包裹完整的loss计算流程，确保所有待优化变量的计算步骤都被记录。
• 避免直接在minimize中手动创建新磁带，而是提前用磁带记录计算过程后再传入，或直接拆分compute_gradients和apply_gradients步骤（更灵活）。

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