TensorFlow 1.13.1中GradientTape计算梯度触发TypeError求助

我之前在TF1.x版本里也遇到过一模一样的问题，本质是这个版本的GradientTape在图模式下，对tf.map_fn（以及Keras RNN内部依赖的循环/子图操作）的梯度追踪存在实现缺陷——这类操作会生成独立的子图结构，早期的GradientTape没法正确捕捉它们和输入张量的依赖关系，导致梯度计算返回None，最后在sess.run()时触发那个烦人的类型错误。而且你已经排除了梯度为0的情况，这就坐实了是版本兼容的问题。

下面给你几个针对性的解决方案，从简单到复杂：

一、替换tf.map_fn为向量化操作（简化场景首选） #

你的第一个示例里，tf.map_fn(lambda a: a**2, x)完全是多余的，直接用TensorFlow的向量化元素操作就能替代，这类操作能被GradientTape完美追踪：

import tensorflow as tf
tf.reset_default_graph()
x = tf.constant([1., 2., 3.])
with tf.GradientTape(persistent=True) as gg:
    gg.watch(x)
    f1 = x ** 2  # 把map_fn换成向量化操作
    f2 = x*x
# 计算梯度
d_fx1 = gg.gradient(f1, x)
d_fx2 = gg.gradient(f2, x)
del gg
with tf.Session() as sess:
    d1, d2 = sess.run((d_fx1, d_fx2))
    print(d1, d2)  # 输出 [2. 4. 6.] [2. 4. 6.]

二、开启Eager模式使用GradientTape（兼容TF1.x的通用方案） #

TF1.x里其实支持启用Eager Execution，开启后GradientTape对循环类操作的追踪会靠谱很多，不管是map_fn还是RNN都能正常计算梯度：

简化示例（Eager模式） #

import tensorflow as tf
tf.enable_eager_execution()  # 开启Eager模式，不用再写Session了
x = tf.constant([1., 2., 3.])
with tf.GradientTape(persistent=True) as gg:
    gg.watch(x)
    f1 = tf.map_fn(lambda a: a**2, x)
    f2 = x*x
d_fx1 = gg.gradient(f1, x)
d_fx2 = gg.gradient(f2, x)
del gg
print(d_fx1.numpy(), d_fx2.numpy())  # 直接输出结果

RNN雅可比矩阵计算（Eager模式） #

import tensorflow as tf
from tensorflow.keras.layers import RNN, GRUCell
tf.enable_eager_execution()

# 定义变量尺寸
inp_dim = 2
num_units = 50
batch_size = 100
timesteps = 10

inputs = tf.ones(shape=(timesteps, batch_size, inp_dim))
with tf.GradientTape() as g:
    g.watch(inputs)
    cells = [GRUCell(num_units), GRUCell(num_units)]
    rnn = RNN(cells, time_major=True, return_sequences=True)
    f = rnn(inputs)
d_fx = g.batch_jacobian(f, inputs)

print(d_fx.shape)  # 输出 (10, 100, 50, 10, 2)，完全符合雅可比矩阵的维度要求

三、改用TF1.x原生tf.gradients函数（必须保持图模式的方案） #

如果你的项目必须在图模式下运行，那可以放弃GradientTape，改用TF1.x原生的tf.gradients函数，它对map_fn和RNN的支持更稳定：

简化示例（tf.gradients版本） #

import tensorflow as tf
tf.reset_default_graph()
x = tf.constant([1., 2., 3.])
f1 = tf.map_fn(lambda a: a**2, x)
f2 = x*x

# tf.gradients返回梯度列表，取第一个元素就是对应输入的梯度
d_fx1 = tf.gradients(f1, x)[0]
d_fx2 = tf.gradients(f2, x)[0]

with tf.Session() as sess:
    d1, d2 = sess.run((d_fx1, d_fx2))
    print(d1, d2)  # 输出正常的梯度结果

RNN雅可比矩阵计算（图模式手动实现） #

TF1.x里没有现成的batch_jacobian，需要手动对每个输出维度计算梯度，这里给你实现一个批量雅可比的版本（每个样本的输出对自身输入的梯度）：

import tensorflow as tf
from tensorflow.keras.layers import RNN, GRUCell

inp_dim = 2
num_units = 50
batch_size = 100
timesteps = 10

tf.reset_default_graph()
inputs = tf.placeholder(tf.float32, shape=(timesteps, batch_size, inp_dim))

# 构建RNN模型
cells = [GRUCell(num_units), GRUCell(num_units)]
rnn = RNN(cells, time_major=True, return_sequences=True)
f = rnn(inputs)

# 手动计算批量雅可比：每个样本的输出对自身输入的梯度
jacobian_list = []
for b in range(batch_size):
    # 对每个样本的输出计算梯度
    sample_output = f[:, b, :]
    sample_input = inputs[:, b, :]
    grad = tf.gradients(sample_output, sample_input)[0]
    jacobian_list.append(grad)

# 整理形状为 (batch_size, timesteps, num_units, timesteps, inp_dim)
jacobian = tf.stack(jacobian_list)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    grads = sess.run(jacobian, feed_dict={inputs: tf.ones((timesteps, batch_size, inp_dim)).numpy()})
    print(grads.shape)

这个版本的计算量比全雅可比小很多，适合大多数业务场景。

终极建议：升级到TensorFlow 2.x #

如果你的环境允许，升级到TF2.x是一劳永逸的解决方案。TF2.x默认启用Eager模式，GradientTape对各种操作的梯度追踪都非常完善，不管是map_fn还是RNN，batch_jacobian用起来也毫无障碍，完全不会出现这类TF1.x的兼容性bug。

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