PyTorch高阶玩法：用torch.autograd.grad的create_graph参数计算模型二阶导（Hessian矩阵入门）

在深度学习的进阶应用中，我们常常需要超越一阶导数的计算能力。想象一下这样的场景：当你需要实现牛顿法优化神经网络参数时，或是构建物理信息神经网络（PINN）来求解微分方程时，二阶导数矩阵（Hessian矩阵）的计算就变得至关重要。PyTorch的 torch.autograd.grad() 函数中的 create_graph 参数，正是打开这扇高阶导数计算之门的钥匙。

本文将带你深入探索如何利用 create_graph=True 参数，将一阶梯度本身变为可微张量，进而实现二阶甚至更高阶导数的计算。不同于基础教程，我们会从实际应用场景出发，通过完整的代码示例展示这一技术在模型优化、元学习等前沿领域的具体应用。无论你是正在研究高级优化算法的工程师，还是需要精确计算模型敏感性的研究员，这些技巧都将为你的工具箱增添强大武器。

1. 理解autograd的计算图机制

PyTorch的自动微分系统（autograd）是其核心优势之一。要掌握高阶导数计算，首先需要深入理解计算图（Computational Graph）的工作机制。计算图是PyTorch记录运算过程的动态数据结构，它会在我们执行张量运算时自动构建。

1.1 计算图的基本概念

当我们对一个张量设置 requires_grad=True 时，PyTorch会跟踪所有与之相关的运算：

import torch

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 + 3 * x  # 计算图开始构建

在这个简单的例子中，PyTorch会自动记录从 x 到 y 的计算路径。当我们调用 y.backward() 时，系统会沿着这个路径反向传播梯度。

1.2 叶子节点与中间节点

理解叶子节点（leaf nodes）和中间节点（intermediate nodes）的区别至关重要：

  • 叶子节点 ：直接创建的张量（如我们的 x ），其 is_leaf 属性为True
  • 中间节点 ：通过运算产生的张量（如 y ），其 is_leaf 属性为False

默认情况下，PyTorch只会保留叶子节点的梯度。要获取中间节点的梯度，我们需要特别处理：

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2  # 中间节点
y.retain_grad()  # 显式要求保留y的梯度
z = y.mean()
z.backward()

print(f"x的梯度: {x.grad}")  # 输出: tensor([4.])
print(f"y的梯度: {y.grad}")  # 输出: tensor([0.5000])

1.3 torch.autograd.grad与backward的区别

PyTorch提供了两种计算梯度的方法：

方法 特点 适用场景
tensor.backward() 计算所有叶子节点的梯度并累加到.grad属性 常规训练场景
torch.autograd.grad() 精确计算指定输出对指定输入的梯度，返回梯度值 需要精确控制梯度计算的高级场景

torch.autograd.grad() 的基本用法如下：

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
dy_dx = torch.autograd.grad(outputs=y, inputs=x)[0]
print(dy_dx)  # 输出: tensor([4.])

这种方法的优势在于可以精确控制计算哪个变量对哪个变量的导数，而不影响其他变量的梯度状态。

2. create_graph参数的核心原理

create_graph 参数是 torch.autograd.grad() 函数中最强大但最容易被忽视的选项。当设置为True时，它会让返回的梯度本身成为可微张量，这是计算高阶导数的关键。

2.1 一阶导数与二阶导数的关系

在数学上，二阶导数是"导数的导数"。同样，在PyTorch中，我们可以：

  1. 计算一阶导数（保持计算图）
  2. 对一阶导数再求导，得到二阶导数
x = torch.tensor([3.0], requires_grad=True)
y = x ** 3 + 2 * x ** 2

# 计算一阶导数，并保留计算图
dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]

# 对一阶导数再求导，得到二阶导数
d2y_dx2 = torch.autograd.grad(dy_dx, x)[0]

print(f"一阶导数: {dy_dx.item()}")  # 3*3² + 2*2*3 = 27 + 12 = 39
print(f"二阶导数: {d2y_dx2.item()}")  # 6*3 + 4 = 22

2.2 create_graph的底层机制

当 create_graph=True 时，PyTorch会：

  1. 保留计算一阶导数所需的所有中间结果
  2. 使返回的梯度张量具有 requires_grad=True 属性
  3. 允许对这些梯度张量进一步求导

我们可以验证这一点：

x = torch.tensor([1.0], requires_grad=True)
y = torch.exp(x)

# 不使用create_graph
grad1 = torch.autograd.grad(y, x)[0]
print(grad1.requires_grad)  # False

# 使用create_graph
grad2 = torch.autograd.grad(y, x, create_graph=True)[0]
print(grad2.requires_grad)  # True

2.3 计算图的内存管理

使用 create_graph=True 会显著增加内存消耗，因为它需要保留更多的中间结果。在实际应用中，我们需要注意：

  • 及时释放不再需要的计算图
  • 合理使用 retain_graph 参数
  • 在循环中谨慎管理内存
x = torch.tensor([2.0], requires_grad=True)

for _ in range(3):
    y = x ** 2
    dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
    print(dy_dx)
    
    # 手动释放计算图
    dy_dx.backward(retain_graph=True)
    x.grad.zero_()  # 清除累积的梯度

3. 实战：Hessian矩阵计算

Hessian矩阵是二阶导数在多维情况下的推广，它包含了函数所有可能的二阶偏导数。在深度学习中，Hessian矩阵在优化算法、不确定性估计等领域有重要应用。

3.1 单变量函数的二阶导数

我们先从简单的单变量函数开始：

def compute_second_derivative(f, x):
    """计算标量函数f在x处的二阶导数"""
    x_tensor = torch.tensor([x], requires_grad=True, dtype=torch.float32)
    y = f(x_tensor)
    
    # 计算一阶导数并保留计算图
    dy_dx = torch.autograd.grad(y, x_tensor, create_graph=True)[0]
    
    # 计算二阶导数
    d2y_dx2 = torch.autograd.grad(dy_dx, x_tensor)[0]
    
    return d2y_dx2.item()

# 测试函数 f(x) = x^3 + sin(x)
def test_func(x):
    return x ** 3 + torch.sin(x)

x_value = 2.0
second_deriv = compute_second_derivative(test_func, x_value)
print(f"在x={x_value}处的二阶导数为: {second_deriv:.4f}")

3.2 多变量函数的Hessian矩阵

对于多变量函数，Hessian矩阵的计算稍微复杂一些。我们需要计算每个变量对所有变量的二阶偏导数：

def compute_hessian(f, variables):
    """计算函数f关于variables的Hessian矩阵"""
    # 计算一阶梯度
    grads = torch.autograd.grad(f, variables, create_graph=True, allow_unused=True)
    
    # 初始化Hessian矩阵
    hessian = []
    for grad in grads:
        if grad is None:
            hessian.append([torch.zeros_like(var) for var in variables])
            continue
            
        # 对每个梯度分量计算关于所有变量的导数
        grad_grad = torch.autograd.grad(grad, variables, retain_graph=True, allow_unused=True)
        grad_grad = [g if g is not None else torch.zeros_like(v) for g, v in zip(grad_grad, variables)]
        hessian.append(grad_grad)
    
    return hessian

# 示例：计算二元函数的Hessian
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x[0] ** 3 + x[1] ** 2 + x[0] * x[1]

hessian = compute_hessian(y, x)
print("Hessian矩阵:")
for row in hessian:
    print([h.item() for h in row])

3.3 高效计算Hessian-向量积

在实际应用中，我们通常不需要完整的Hessian矩阵，而是需要Hessian-向量积。这种方法更高效：

def hessian_vector_product(f, variables, vector):
    """计算Hessian-向量积"""
    # 计算一阶梯度
    grads = torch.autograd.grad(f, variables, create_graph=True)
    
    # 计算梯度与向量的点积
    grad_vector = sum(torch.sum(g * v) for g, v in zip(grads, vector))
    
    # 计算Hessian-向量积
    hvp = torch.autograd.grad(grad_vector, variables)
    
    return hvp

# 示例使用
x = torch.tensor([1.0, 2.0], requires_grad=True)
v = torch.tensor([0.5, 0.5])  # 任意向量
y = x[0] ** 3 + x[1] ** 2

hvp = hessian_vector_product(y, x, v)
print(f"Hessian-向量积: {[h.item() for h in hvp]}")

4. 高阶导数在前沿场景中的应用

掌握了高阶导数计算技术后，我们可以将其应用于多个前沿领域。以下是几个典型的应用场景。

4.1 牛顿法优化

牛顿法利用Hessian矩阵进行二阶优化，相比一阶梯度下降法，它能提供更精确的更新方向：

def newton_method(f, initial_x, lr=0.01, steps=10):
    """使用牛顿法优化函数f"""
    x = torch.tensor(initial_x, requires_grad=True)
    
    for i in range(steps):
        # 计算函数值
        y = f(x)
        
        # 计算一阶梯度
        dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
        
        # 计算Hessian
        hessian = []
        for grad in dy_dx:
            h = torch.autograd.grad(grad, x, retain_graph=True)[0]
            hessian.append(h)
        hessian = torch.stack(hessian)
        
        # 牛顿法更新: x = x - H^{-1} * grad
        delta = torch.linalg.solve(hessian, dy_dx)
        x = x - lr * delta
        
        print(f"Step {i+1}: x = {x.detach().numpy()}, f(x) = {y.item():.4f}")
    
    return x

# 测试函数
def test_function(x):
    return (x[0] - 1) ** 4 + (x[1] - 2) ** 2 + x[0] * x[1]

initial_x = [0.0, 0.0]
optimal_x = newton_method(test_function, initial_x)

4.2 元学习中的二阶优化

在模型无关的元学习（MAML）等算法中，二阶导数计算对于获得准确的元梯度至关重要：

def maml_step(model, loss_fn, tasks, inner_lr=0.1, meta_lr=0.01):
    """简化的MAML元更新步骤"""
    meta_grads = None
    
    for task in tasks:
        # 内循环适应
        x, y = task
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        
        # 计算内循环梯度并创建计算图
        grads = torch.autograd.grad(loss, model.parameters(), create_graph=True)
        
        # 创建快速权重（临时参数）
        fast_weights = [p - inner_lr * g for p, g in zip(model.parameters(), grads)]
        
        # 在新任务上评估
        x_val, y_val = task.get_validation()
        y_pred_val = model(x_val, fast_weights)
        val_loss = loss_fn(y_pred_val, y_val)
        
        # 计算元梯度（需要二阶导数）
        if meta_grads is None:
            meta_grads = torch.autograd.grad(val_loss, model.parameters())
        else:
            new_grads = torch.autograd.grad(val_loss, model.parameters())
            meta_grads = [mg + ng for mg, ng in zip(meta_grads, new_grads)]
    
    # 元更新
    for p, g in zip(model.parameters(), meta_grads):
        p.data -= meta_lr * g / len(tasks)

4.3 物理信息神经网络（PINN）

在求解微分方程时，我们需要计算函数的高阶导数：

class PINN(nn.Module):
    """物理信息神经网络"""
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 20),
            nn.Tanh(),
            nn.Linear(20, 20),
            nn.Tanh(),
            nn.Linear(20, 1)
        )
    
    def forward(self, x):
        return self.net(x)
    
    def compute_pde_loss(self, x):
        """计算PDE残差"""
        x.requires_grad_(True)
        u = self(x)
        
        # 计算一阶导数
        du_dx = torch.autograd.grad(u, x, create_graph=True, grad_outputs=torch.ones_like(u))[0]
        
        # 计算二阶导数
        d2u_dx2 = torch.autograd.grad(du_dx, x, create_graph=True)[0]
        
        # 示例PDE: u'' + u = 0
        pde_residual = d2u_dx2 + u
        return torch.mean(pde_residual ** 2)

4.4 对抗样本生成

在生成对抗样本时，高阶导数可以帮助我们找到更有效的扰动方向：

def second_order_attack(model, input, target, epsilon=0.1, steps=10):
    """使用二阶信息的对抗攻击"""
    perturbed = input.clone().requires_grad_(True)
    
    for _ in range(steps):
        output = model(perturbed)
        loss = F.cross_entropy(output, target)
        
        # 计算一阶梯度
        grad = torch.autograd.grad(loss, perturbed, create_graph=True)[0]
        
        # 计算Hessian-向量积（使用一阶梯度作为向量）
        hvp = torch.autograd.grad(torch.sum(grad * grad), perturbed)[0]
        
        # 二阶更新
        perturbation = grad + 0.5 * epsilon * hvp
        perturbation = epsilon * perturbation / perturbation.norm()
        
        perturbed = perturbed.detach() + perturbation
        perturbed = torch.clamp(perturbed, 0, 1).requires_grad_(True)
    
    return perturbed

5. 性能优化与调试技巧

高阶导数计算在带来强大功能的同时，也带来了性能挑战。下面是一些优化和调试的技巧。

5.1 减少内存消耗的策略

高阶导数计算会显著增加内存使用，以下策略可以帮助缓解：

  • 使用checkpointing ：只保留必要的中间结果
  • 及时释放计算图 ：在不需要时及时调用 retain_graph=False
  • 分批计算 ：对大矩阵分块计算
from torch.utils.checkpoint import checkpoint

def memory_efficient_hessian(f, x):
    """内存高效的Hessian计算"""
    # 使用checkpointing计算一阶梯度
    def get_grad(x):
        y = f(x)
        return torch.autograd.grad(y, x, create_graph=True)[0]
    
    grad = checkpoint(get_grad, x)
    
    # 分块计算Hessian
    hessian = []
    for i in range(len(x)):
        # 仅计算对角线元素
        h_i = torch.autograd.grad(grad[i], x, retain_graph=i<len(x)-1)[0]
        hessian.append(h_i[i].item())  # 只存储对角线元素
    
    return hessian

5.2 常见错误与解决方法

在使用高阶导数时，你可能会遇到以下问题：

错误类型 原因 解决方法
"Trying to backward through the graph a second time" 计算图已被释放 设置 retain_graph=True
"One of the differentiated Tensors appears to not have been used in the graph" 输入变量未参与计算 检查计算路径或设置 allow_unused=True
"CUDA out of memory" 高阶导数占用内存过多 减少批量大小或使用checkpointing
"Gradients do not seem to be correct" 计算图构建错误 验证中间结果的 requires_grad 属性

5.3 数值梯度验证

为确保自动微分计算的正确性，可以使用数值梯度进行验证：

def numerical_second_derivative(f, x, h=1e-5):
    """计算数值二阶导数"""
    return (f(x + h) - 2 * f(x) + f(x - h)) / (h ** 2)

# 比较自动微分和数值微分结果
x_test = 2.0
f = lambda x: x ** 3 + torch.sin(x)

auto_deriv = compute_second_derivative(f, x_test)
num_deriv = numerical_second_derivative(lambda x: f(torch.tensor(x)).item(), x_test)

print(f"自动微分结果: {auto_deriv:.6f}")
print(f"数值微分结果: {num_deriv:.6f}")
print(f"相对误差: {abs(auto_deriv - num_deriv) / auto_deriv:.2%}")

5.4 高阶导数可视化

可视化可以帮助理解高阶导数的行为：

import matplotlib.pyplot as plt
import numpy as np

def plot_derivatives(f, x_range=(-3, 3)):
    """绘制函数及其导数"""
    x_vals = np.linspace(*x_range, 100)
    y_vals = []
    dy_vals = []
    d2y_vals = []
    
    for x in x_vals:
        x_tensor = torch.tensor([x], requires_grad=True, dtype=torch.float32)
        y = f(x_tensor)
        
        # 一阶导数
        dy = torch.autograd.grad(y, x_tensor, create_graph=True)[0]
        
        # 二阶导数
        d2y = torch.autograd.grad(dy, x_tensor)[0]
        
        y_vals.append(y.item())
        dy_vals.append(dy.item())
        d2y_vals.append(d2y.item())
    
    plt.figure(figsize=(10, 6))
    plt.plot(x_vals, y_vals, label="f(x)")
    plt.plot(x_vals, dy_vals, label="f'(x)")
    plt.plot(x_vals, d2y_vals, label="f''(x)")
    plt.legend()
    plt.grid()
    plt.title("函数及其导数可视化")
    plt.show()

# 示例函数
def example_func(x):
    return torch.sin(x) + 0.5 * x ** 2

plot_derivatives(example_func)
# 深度学习
Logo
https://edu.csdn.net/learn/39067/627173?utm_source=2019755004

汇聚全球AI编程工具，助力开发者即刻编程。

更多推荐

cover

一个 API 入口调用多个大模型：AiiOnly客户端 + CC Switch 打通 Codex AI编程流程

avatar AI编程社区
cover

VSCode Claude Code + MiMo 中转免登录完整配置教程（解决原生登录弹窗 + 模型报错）

avatar AI编程社区
cover

从零开始做一个属于自己的 Skill（保姆级教程）

avatar AI编程社区