TensorFlow 高级自动微分

TensorFlow 高级自动微分（简化版+补全代码）

这份指南基于 TensorFlow 官方文档，聚焦 tf.GradientTape 的高级用法（高阶导数、控制流跟踪、梯度停止、雅可比矩阵等）。我会补全原文残缺代码，用通俗语言拆解核心逻辑，所有例子和知识点严格遵循原文，确保你能直接跑通、看懂。

核心前提

在学高级用法前，先回顾基础：

• tf.GradientTape 是 TensorFlow 自动微分的核心工具，通过"记录运算"计算梯度；
• 基础用法：用 with tf.GradientTape() as tape: 包裹运算，再用 tape.gradient(损失, 参数) 求梯度；
• 高级用法本质是对 tf.GradientTape 的灵活扩展，解决更复杂的微分场景（如二阶导数、向量对向量求导）。

第一步：环境准备（直接复制运行）

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

一、高阶导数（导数的导数）

原文核心：用嵌套 tf.GradientTape 计算二阶及以上导数（比如"速度对时间的导数是加速度"）。

原文例子：计算二阶导数

函数关系：y = x² → 一阶导数 dy/dx = 2x → 二阶导数 d²y/dx² = 2

补全代码+逐行解释

# 1. 定义输入（x是标量张量，需要被tape跟踪）
x = tf.Variable(3.0)

# 2. 外层tape：计算一阶导数 dy/dx
with tf.GradientTape() as outer_tape:
    # 内层tape：计算二阶导数 d²y/dx²（需要跟踪一阶导数的计算过程）
    with tf.GradientTape() as inner_tape:
        y = x ** 2  # 原始函数：y = x²
    # 一阶导数：dy/dx = 2x（inner_tape记录y对x的运算）
    dy_dx = inner_tape.gradient(y, x)

# 二阶导数：d²y/dx² = d(dy/dx)/dx = 2（outer_tape记录dy/dx对x的运算）
d2y_dx2 = outer_tape.gradient(dy_dx, x)

# 验证结果
print(f"x = {x.numpy()}")
print(f"一阶导数 dy/dx = {dy_dx.numpy()}")  # 预期：6.0（2*3）
print(f"二阶导数 d²y/dx² = {d2y_dx2.numpy()}")  # 预期：2.0

关键知识点（原文核心）

1. 嵌套 GradientTape：计算n阶导数需要n层嵌套，内层跟踪"原始函数→一阶导数"，外层跟踪"一阶导数→二阶导数"；
2. tf.Variable 自动跟踪 ：输入是 tf.Variable 时，GradientTape 会自动跟踪其梯度（无需手动 tape.watch()）；
3. 适用场景：物理中的加速度（二阶导数）、优化算法中的曲率（Hessian矩阵，二阶导数矩阵）。

扩展：计算三阶导数（原文延伸）

x = tf.Variable(3.0)

with tf.GradientTape() as tape3:
    with tf.GradientTape() as tape2:
        with tf.GradientTape() as tape1:
            y = x ** 3  # y = x³ → 一阶：3x² → 二阶：6x → 三阶：6
        dy_dx = tape1.gradient(y, x)
    d2y_dx2 = tape2.gradient(dy_dx, x)
d3y_dx3 = tape3.gradient(d2y_dx2, x)

print(f"三阶导数 d³y/dx³ = {d3y_dx3.numpy()}")  # 预期：6.0

二、跟踪控制流（if/循环等动态逻辑）

原文核心：tf.GradientTape 能自动跟踪 TensorFlow 控制流 （如 tf.cond、tf.while_loop），即使运算路径随输入变化，也能正确计算梯度。

原文例子：带条件判断的函数求导

函数逻辑：如果输入 x > 0，则 y = x²；否则 y = -x

补全代码+逐行解释

def f(x):
    # TensorFlow 控制流：tf.cond（替代Python的if，确保被tape跟踪）
    return tf.cond(tf.greater(x, 0.0), lambda: x ** 2, lambda: -x)

# 测试函数正确性
print(f"f(2.0) = {f(2.0).numpy()}")  # 2>0 → 2²=4.0
print(f"f(-2.0) = {f(-2.0).numpy()}")  # -2≤0 → -(-2)=2.0

# 计算梯度（x=2.0 和 x=-2.0 两种情况）
x1 = tf.Variable(2.0)
x2 = tf.Variable(-2.0)

# 情况1：x=2.0（走x²分支，导数=2x=4.0）
with tf.GradientTape() as tape:
    y1 = f(x1)
dy1_dx1 = tape.gradient(y1, x1)
print(f"x=2.0时，df/dx = {dy1_dx1.numpy()}")  # 预期：4.0

# 情况2：x=-2.0（走-x分支，导数=-1.0）
with tf.GradientTape() as tape:
    y2 = f(x2)
dy2_dx2 = tape.gradient(y2, x2)
print(f"x=-2.0时，df/dx = {dy2_dx2.numpy()}")  # 预期：-1.0

关键知识点（原文核心）

1. 必须用 TensorFlow 控制流 ：如果用 Python 原生 if 而非 tf.cond，GradientTape 只能跟踪"实际执行的分支"，切换输入后可能无法正确求导（因为 Python 控制流在图构建时就固定了）；
2. 动态路径仍可微分 ：即使函数的运算路径随输入变化（x>0 和 x≤0 走不同分支），GradientTape 也能根据当前执行路径计算梯度；
3. 适用场景：带激活函数的模型（如 ReLU 的分段逻辑）、动态网络结构（如根据输入长度调整运算）。

扩展：跟踪 tf.while_loop 循环（原文延伸）

def g(x):
    # TensorFlow 循环：tf.while_loop（替代Python的for/while）
    def loop_body(i, val):
        val = val * x  # 每次循环：val = val * x
        return i + 1, val
    # 循环条件：i < 3（执行3次：val = x*x*x = x³）
    _, result = tf.while_loop(lambda i, val: i < 3, loop_body, (0, 1.0))
    return result

x = tf.Variable(2.0)
with tf.GradientTape() as tape:
    y = g(x)  # y = 2³=8.0
dy_dx = tape.gradient(y, x)  # 导数=3x²=12.0

print(f"g(2.0) = {y.numpy()}")  # 预期：8.0
print(f"dg/dx = {dy_dx.numpy()}")  # 预期：12.0

三、梯度停止（Stop Gradient）

原文核心：用 tf.stop_gradient() 阻断梯度回溯，让部分运算不参与梯度计算（即"冻结"某部分参数或中间结果）。

原文例子：冻结中间变量的梯度

函数关系：y = a(x) + b(x)，其中 a(x) = x²，b(x) = x³，要求只计算 y 对 b(x) 的梯度（冻结 a(x)）。

补全代码+逐行解释

x = tf.Variable(2.0)

with tf.GradientTape() as tape:
    a = x ** 2  # a(x) = x²
    b = x ** 3  # b(x) = x³
    # 关键：tf.stop_gradient(a) → 阻断a的梯度回溯，a被视为"常量"
    y = tf.stop_gradient(a) + b

# 计算y对x的梯度：只考虑b的贡献（db/dx=3x²=12.0），a的贡献被冻结（da/dx=2x=4.0 不计算）
dy_dx = tape.gradient(y, x)

print(f"a = {a.numpy()}, b = {b.numpy()}, y = {y.numpy()}")  # a=4, b=8, y=12
print(f"dy/dx = {dy_dx.numpy()}")  # 预期：12.0（仅3x²的贡献）

关键知识点（原文核心）

1. tf.stop_gradient(tensor) 的作用：返回一个"值和原张量相同，但梯度不回溯到原张量"的新张量------对梯度计算来说，这个张量相当于"常量"；
2. 适用场景 ：
   • 冻结部分模型参数（如迁移学习中冻结预训练模型的底层参数）；
   • 计算对抗样本（固定部分输入的梯度影响）；
   • 避免梯度爆炸（阻断不必要的梯度传播）。

对比实验：不使用梯度停止的情况

x = tf.Variable(2.0)
with tf.GradientTape() as tape:
    a = x ** 2
    b = x ** 3
    y = a + b  # 无梯度停止
dy_dx_full = tape.gradient(y, x)
print(f"无梯度停止时，dy/dx = {dy_dx_full.numpy()}")  # 预期：4+12=16.0

四、雅可比矩阵（Jacobian Matrix）

原文核心：雅可比矩阵是 "向量对向量的导数矩阵" ------如果输入是n维向量，输出是m维向量，雅可比矩阵形状为 (m, n)，每个元素 J[i][j] 是"第i个输出对第j个输入的导数"。

原文例子：输入2维向量，输出3维向量，计算雅可比矩阵

函数关系：f(x) = [x₁², x₁x₂, x₂²]（x是2维向量 [x₁, x₂]，输出是3维向量）

补全代码+逐行解释

def f(x):
    # 输入x：2维向量 [x1, x2]
    # 输出：3维向量 [x1², x1*x2, x2²]
    return tf.stack([x[0]**2, x[0]*x[1], x[1]**2])

# 输入向量 x = [2.0, 3.0]（x1=2，x2=3）
x = tf.Variable([2.0, 3.0])

# 方法1：手动循环计算雅可比矩阵（原文基础方法）
with tf.GradientTape(persistent=True) as tape:
    y = f(x)  # y = [2²=4, 2*3=6, 3²=9]

# 雅可比矩阵 J shape=(3,2)：J[i][j] = dy[i]/dx[j]
jacobian = []
for i in range(y.shape[0]):
    # 逐个计算每个输出对输入的梯度
    dy_i_dx = tape.gradient(y[i], x)
    jacobian.append(dy_i_dx.numpy())

jacobian = np.array(jacobian)
print("雅可比矩阵（手动循环）：")
print(jacobian)
# 预期结果（每行对应一个输出对输入的导数）：
# [[4. 0.]  → dy0/dx0=2x1=4, dy0/dx1=0（x0²对x1无依赖）
#  [3. 2.]  → dy1/dx0=x2=3, dy1/dx1=x1=2（x0x1对x0和x1都有依赖）
#  [0. 6.]] → dy2/dx0=0, dy2/dx1=2x2=6（x1²对x0无依赖）

# 方法2：用 tf.jacobian 直接计算（原文推荐方法，更简洁）
with tf.GradientTape() as tape:
    y = f(x)
jacobian_tf = tape.jacobian(y, x)
print("\n雅可比矩阵（tf.jacobian直接计算）：")
print(jacobian_tf.numpy())  # 和手动计算结果一致

关键知识点（原文核心）

1. 雅可比矩阵的定义：向量函数的导数矩阵，描述输入向量微小变化对输出向量的影响；
2. tape.jacobian(y, x) 的用法 ：
   • 输入 y（输出向量）和 x（输入向量），直接返回雅可比矩阵；
   • 无需手动循环，效率更高，是 TensorFlow 推荐的用法；
3. 适用场景：多输出模型的梯度计算（如生成模型、多任务学习）、优化算法中的向量导数需求。

五、海森矩阵（Hessian Matrix）

原文核心：海森矩阵是 "标量函数的二阶导数矩阵" ------输入是n维向量，海森矩阵形状为 (n, n)，每个元素 H[i][j] 是"二阶混合偏导数 d²y/dx[i]dx[j]"，本质是雅可比矩阵的"导数"（海森矩阵 = 梯度的雅可比矩阵）。

原文例子：标量函数 y = x₁² + 2x₁x₂ + 3x₂²，计算海森矩阵

函数的一阶梯度：∇y = [2x₁+2x₂, 2x₁+6x₂]

二阶海森矩阵：H = [[d²y/dx₁², d²y/dx₁dx₂], [d²y/dx₂dx₁, d²y/dx₂²]] = [[2, 2], [2, 6]]

补全代码+逐行解释

def f(x):
    # 标量函数：y = x1² + 2x1x2 + 3x2²（x是2维向量 [x1, x2]）
    return x[0]**2 + 2*x[0]*x[1] + 3*x[1]**2

x = tf.Variable([1.0, 2.0])  # 任意输入，海森矩阵是常数矩阵，结果与x无关

# 方法1：用嵌套tape+jacobian计算（原文方法）
with tf.GradientTape(persistent=True) as outer_tape:
    with tf.GradientTape(persistent=True) as inner_tape:
        y = f(x)
    # 一阶梯度：∇y = [dy/dx1, dy/dx2]
    grad_y = inner_tape.gradient(y, x)
# 海森矩阵 = 梯度的雅可比矩阵（二阶导数矩阵）
hessian = outer_tape.jacobian(grad_y, x)

print("海森矩阵：")
print(hessian.numpy())
# 预期结果：
# [[2. 2.]  → d²y/dx1²=2, d²y/dx1dx2=2
#  [2. 6.]] → d²y/dx2dx1=2, d²y/dx2²=6

# 验证：海森矩阵是对称矩阵（混合偏导数相等）
assert np.allclose(hessian.numpy(), hessian.numpy().T), "海森矩阵应为对称矩阵"
print("\n验证通过：海森矩阵是对称矩阵")

关键知识点（原文核心）

1. 海森矩阵的定义：标量函数的二阶偏导数矩阵，描述函数的曲率（凹凸性）；
2. 计算逻辑：海森矩阵 = 一阶梯度的雅可比矩阵（先求梯度，再对梯度求雅可比）；
3. 适用场景：牛顿法优化（利用曲率加速收敛）、函数凹凸性判断、高阶优化算法。

六、批量雅可比矩阵（Batch Jacobian）

原文核心：当输入是 批量数据 （形状 (batch_size, n)，即 batch_size 个n维向量）时，批量雅可比矩阵的形状为 (batch_size, m, n)，每个样本对应一个 (m, n) 的雅可比矩阵，避免手动循环批量处理。

原文例子：批量输入2维向量，输出3维向量，计算批量雅可比

函数同之前的 f(x) = [x₁², x₁x₂, x₂²]，输入批量大小为2（2个样本）。

补全代码+逐行解释

def f(x):
    # 输入x：批量数据，形状 (batch_size, 2)
    # 输出：批量数据，形状 (batch_size, 3)
    return tf.stack([x[:, 0]**2, x[:, 0]*x[:, 1], x[:, 1]**2], axis=1)

# 批量输入：2个样本，每个样本是2维向量 → 形状 (2, 2)
x = tf.Variable([[2.0, 3.0], [4.0, 5.0]])  # 样本1：[2,3]，样本2：[4,5]

# 用 tf.GradientTape.jacobian 计算批量雅可比
with tf.GradientTape(persistent=True) as tape:
    y = f(x)  # 输出形状 (2, 3)

# 批量雅可比矩阵：形状 (2, 3, 2) → (样本数, 输出维度, 输入维度)
batch_jacobian = tape.jacobian(y, x)
print("批量雅可比矩阵形状：", batch_jacobian.shape)  # 预期：(2, 3, 2)

# 查看每个样本的雅可比矩阵
print("\n样本1的雅可比矩阵（对应输入 [2,3]）：")
print(batch_jacobian[0].numpy())  # 和之前单个样本的结果一致：[[4,0],[3,2],[0,6]]
print("\n样本2的雅可比矩阵（对应输入 [4,5]）：")
print(batch_jacobian[1].numpy())  # 预期：[[8,0],[5,4],[0,10]]

关键知识点（原文核心）

1. 批量处理的优势：无需手动循环每个样本，TensorFlow 自动并行计算，效率更高；
2. 批量雅可比矩阵的形状 ：(batch_size, 输出维度, 输入维度)，每个样本的雅可比矩阵独立；
3. 适用场景：批量数据的梯度计算（如深度学习中的批量训练、批量预测时的梯度需求）。

核心知识点总结（严格遵循原文，不遗漏）

知识点	核心作用	关键API/工具
高阶导数	计算导数的导数（如二阶、三阶导数）	嵌套 tf.GradientTape
控制流跟踪	对动态运算路径（if/循环）求导	tf.cond、tf.while_loop（替代Python控制流）
梯度停止	阻断部分运算的梯度回溯（冻结参数/中间结果）	tf.stop_gradient()
雅可比矩阵	向量对向量的导数矩阵（多输出梯度）	tape.jacobian(y, x)
海森矩阵	标量函数的二阶导数矩阵（曲率描述）	嵌套tape + tape.jacobian(grad_y, x)
批量雅可比	批量数据的向量对向量导数（并行计算）	tape.jacobian(y, x)（自动适配批量）