TensorFlow tf.GradientTape 怎么用？自动微分和常见陷阱详解

tf.GradientTape 是 TF 2.x 的自动微分工具：在 with tf.GradientTape() as tape: 里执行的前向运算会被记录下来，之后调用 tape.gradient(target, sources) 就能自动算出梯度。整个机制就是链式法则——从输出往回走，每一步操作都知道怎么求导，一路乘回来。

python
# 最核心的训练步骤模板
with tf.GradientTape() as tape:
    predictions = model(x_batch, training=True)  # 前向传播
    loss = loss_fn(y_batch, predictions)          # 算损失

gradients = tape.gradient(loss, model.trainable_variables)  # 反向传播
optimizer.apply_gradients(zip(gradients, model.trainable_variables))  # 更新参数

三个最容易踩的坑

1. Tape 用完即废：默认 tape.gradient() 只能调一次，第二次调返回 None。需要多次求梯度就加 persistent=True，用完记得 del tape 释放资源。

2. 只监控 Variable：Tape 默认只追踪 tf.Variable。如果你对 tf.constant 求导会得到 None——需要手动 tape.watch(x) 让它监控。常见场景：对输入 x 求梯度（如对抗样本、显著性图）时，x 是 constant 不是 Variable。

3. 梯度为 None：除了上面两种情况，还有一种隐蔽原因——计算路径断开了。比如 y = x * tf.stop_gradient(z)，x 到 y 的梯度被 stop_gradient 截断了。另外，Variable(trainable=False) 也不会被追踪。

高阶导数：嵌套 Tape

求二阶导需要嵌套两层 Tape：外层记录一阶导的计算过程，内层记录原函数。y = x³ 的一阶导 3x²，二阶导 6x：

python
x = tf.Variable(3.0)
with tf.GradientTape() as tape2:
    with tf.GradientTape() as tape1:
        y = x ** 3
    dy_dx = tape1.gradient(y, x)    # 27.0 (= 3 * 3²)
d2y_dx2 = tape2.gradient(dy_dx, x)  # 18.0 (= 6 * 3)

梯度裁剪

梯度爆炸时用裁剪保命：tf.clip_by_norm(g, max_norm=1.0) 把梯度向量的 L2 范数限制在 1.0 以内。这在 RNN/LSTM 训练中几乎标配——不做裁剪很容易梯度爆炸导致 NaN。

python
gradients = tape.gradient(loss, model.trainable_variables)
gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

追问

GradientTape 和 PyTorch 的 autograd 有什么区别？

PyTorch 的 autograd 是隐式的——只要张量设了 requires_grad=True，所有操作自动记录，不需要手动包 with 块。TF 的 GradientTape 是显式的——必须在 with 块内的操作才会被记录。TF 的设计更省内存（不记录不需要的运算），PyTorch 的设计更方便（少写代码）。实际使用中，TF 训练循环比 PyTorch 多几行，但逻辑等价。

什么时候用 persistent=True？

一个 Tape 对多个目标分别求梯度时。比如 GAN 训练中，判别器的损失对生成器和判别器都需要求梯度；或者一个 loss 对多种参数分组求梯度。但 persistent=True 会保留所有中间结果直到手动删除，显存占用翻倍——不用的时候别开。

tape.gradient 返回 None 怎么排查？

按顺序检查：(1) source 是不是 Variable 或被 watch 了；(2) source 的 trainable 是不是 True；(3) target 到 source 的计算路径有没有被 stop_gradient 截断；(4) 是不是已经调过一次 gradient 了（默认 Tape 只能调一次）；(5) 在 @tf.function 里用 Tape 要确保变量创建在函数外部。