★ FEATURED ARTICLE 【Bug已解决】How to compute Hessian of the loss w.r.t. the parameters in PyTorch using autograd.grad 解决方案问题标签PyTorch、Hessian矩阵、autograd.grad、二阶导数、二阶优化适用版本PyTorch 1.8难度等级⭐⭐⭐⭐⭐高一、问题描述在深度学习中大多数优化算法SGD、Adam只使用一阶梯度信息。但在某些场景下我们需要计算二阶导数Hessian 矩阵二阶优化方法Newtons Method, Trust Region模型不确定性估计Laplace Approximation影响函数Influence Functions剪枝与压缩Optimal Brain Damage损失曲面分析核心问题如何使用 PyTorch 的autograd.grad计算损失函数对模型参数的 Hessian 矩阵Hessian 矩阵的定义对于损失函数 ( L(\theta) )Hessian 矩阵 ( H ) 是一个 ( n \times n \ 的矩阵( n ) 为参数数量[ H_{ij} \frac{\partial^2 L}{\partial \theta_i \partial \theta_j} ]二、错误复现场景尝试用autograd.grad直接计算 Hessianimport torch # 创建参数和损失 x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) y (x ** 2).sum() # 计算一阶梯度 grad torch.autograd.grad(y, x, create_graphTrue)[0] print(f一阶梯度: {grad}) # tensor([2., 4., 6.]) # 尝试直接计算 Hessian错误方式 try: hessian torch.autograd.grad(grad, x)[0] # 看似合理 print(fHessian: {hessian}) except Exception as e: print(f错误: {e})问题上述代码对于简单函数可以工作但对于复杂模型如神经网络直接计算 Hessian 会遇到以下问题内存爆炸Hessian 是 ( n \times n ) 矩阵参数量百万级时无法存储计算效率低需要 ( n ) 次反向传播API 使用不当autograd.grad的参数配置不正确三、根因分析3.1 autograd.grad 的工作原理torch.autograd.grad(outputs, inputs, ...)计算outputs对inputs的梯度。一阶导数: autograd.grad(loss, params) → ∂L/∂θ 二阶导数: autograd.grad(∂L/∂θ, params) → ∂²L/∂θ²关键计算一阶导数时必须设置create_graphTrue这样一阶导数本身也参与计算图才能对其求导。3.2 Hessian 矩阵的计算挑战挑战说明内存复杂度( O(n^2) )n百万参数时需要 TB 级内存计算复杂度( O(n) ) 次反向传播每次 ( O(n) )数值稳定性二阶导数容易数值溢出稀疏性Hessian 通常稀疏但 PyTorch 默认稠密存储3.3 PyTorch 中的 Hessian 计算 APIPyTorch 提供了多种计算 Hessian 的方式API说明适用场景torch.autograd.functional.hessian计算函数的 Hessian小规模函数torch.autograd.functional.jacobian计算 Jacobian 矩阵中间步骤手动autograd.grad循环逐行计算 Hessian灵活控制Hessian-vector product (HVP)计算 Hessian 与向量乘积大规模问题四、解决方案方案一使用torch.autograd.functional.hessian小规模推荐import torch from torch.autograd.functional import hessian # 定义函数 def loss_function(params): 示例: f(x, y) x² 3xy 2y² x, y params[0], params[1] return x**2 3*x*y 2*y**2 # 计算点 point torch.tensor([1.0, 2.0]) # 计算 Hessian H hessian(loss_function, point) print(fHessian 矩阵:\n{H}) # 理论值: [[2, 3], [3, 4]] # ∂²f/∂x² 2, ∂²f/∂x∂y 3, ∂²f/∂y² 4方案二手动使用autograd.grad计算 Hessianimport torch def compute_hessian_manual(loss, params): 手动计算 Hessian 矩阵 原理: 1. 计算一阶梯度create_graphTrue 2. 对一阶梯度的每个分量求导 3. 拼接成 Hessian 矩阵 参数: loss: 标量损失 params: 参数张量扁平化为一维 # 计算一阶梯度 grad torch.autograd.grad( loss, params, create_graphTrue, retain_graphTrue )[0] n grad.numel() hessian torch.zeros(n, n) # 逐行计算 Hessian for i in range(n): # 对第 i 个梯度分量求导 grad2 torch.autograd.grad( grad[i], params, retain_graphTrue, create_graphTrue )[0] hessian[i] grad2 return hessian # 测试 x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) loss (x ** 2).sum() x[0] * x[1] # x0² x1² x2² x0*x1 H compute_hessian_manual(loss, x) print(f参数: {x.data}) print(f损失: {loss.item():.4f}) print(fHessian:\n{H}) # 理论值: [[2, 1, 0], [1, 2, 0], [0, 0, 2]]方案三计算神经网络参数的 Hessianimport torch import torch.nn as nn def compute_nn_hessian(model, loss_fn, inputs, targets): 计算神经网络损失对参数的 Hessian 注意: 仅适用于小模型参数量 1000 # 前向传播 outputs model(inputs) loss loss_fn(outputs, targets) # 收集所有参数 params [p for p in model.parameters() if p.requires_grad] # 扁平化参数 flat_params torch.cat([p.view(-1) for p in params]) # 计算一阶梯度 grad torch.autograd.grad( loss, params, create_graphTrue, retain_graphTrue ) flat_grad torch.cat([g.view(-1) for g in grad]) n flat_grad.numel() hessian torch.zeros(n, n) # 逐行计算 for i in range(n): grad2 torch.autograd.grad( flat_grad[i], params, retain_graphTrue ) flat_grad2 torch.cat([g.view(-1) for g in grad2]) hessian[i] flat_grad2 return hessian # 测试小模型 model nn.Linear(3, 2, biasFalse) # 6 个参数 inputs torch.randn(4, 3) targets torch.tensor([0, 1, 0, 1]) loss_fn nn.CrossEntropyLoss() H compute_nn_hessian(model, loss_fn, inputs, targets) print(f模型参数量: {sum(p.numel() for p in model.parameters())}) print(fHessian 形状: {H.shape}) print(fHessian:\n{H})方案四Hessian-Vector Product (HVP) —— 大规模推荐对于大模型不需要显式计算 Hessian 矩阵只需计算 Hessian 与向量的乘积。import torch import torch.nn as nn def hvp(loss, params, vector): 计算 Hessian-Vector Product: H × v 原理Pearlmutter, 1992: H × v ∇(∇L · v) 即: 先计算梯度和向量的点积再对参数求导 优点: 只需 2 次反向传播不需要存储 Hessian 矩阵 复杂度: O(n) 时间和内存而非 O(n²) # 计算一阶梯度 grad torch.autograd.grad( loss, params, create_graphTrue, retain_graphTrue ) # 梯度与向量点积 grad_v_product sum( (g * v).sum() for g, v in zip(grad, vector) ) # 对点积求导得到 HVP hvp_result torch.autograd.grad( grad_v_product, params, retain_graphTrue ) return hvp_result # 测试 HVP model nn.Linear(10, 5) inputs torch.randn(8, 10) targets torch.randn(8, 5) loss_fn nn.MSELoss() outputs model(inputs) loss loss_fn(outputs, targets) params [p for p in model.parameters() if p.requires_grad] vector [torch.randn_like(p) for p in params] # 随机向量 # 计算 HVP hvp_result hvp(loss, params, vector) for i, (p, h) in enumerate(zip(params, hvp_result)): print(f参数 {i} HVP: shape{h.shape}, norm{h.norm():.4f})方案五使用 HVP 近似 Hessian 的对角线import torch import torch.nn as nn def hessian_diagonal(loss, params, num_iterations1): 使用 HVP 计算 Hessian 的对角线元素 原理: H × e_i H 的第 i 列 对角线: H_ii (H × e_i)_i 对于每个参数使用单位向量计算 HVP diag [] for p in params: # 对每个参数张量逐元素计算对角线 p_diag torch.zeros_like(p) n p.numel() for i in range(n): # 创建单位向量 e_i torch.zeros(n, devicep.device) e_i[i] 1.0 # 重塑为参数形状 e_i e_i.view_as(p) # 构建完整向量其他参数为零 vector [torch.zeros_like(pp) for pp in params] vector[params.index(p)] e_i # 计算 HVP hvp_result hvp(loss, params, vector) # 取对应位置的对角线值 p_diag.view(-1)[i] hvp_result[params.index(p)].view(-1)[i] diag.append(p_diag) return diag # 更高效的对角线估计方法 def hessian_diagonal_efficient(model, loss_fn, inputs, targets): 高效计算 Hessian 对角线一次反向传播 outputs model(inputs) loss loss_fn(outputs, targets) params [p for p in model.parameters() if p.requires_grad] # 计算一阶梯度 grad torch.autograd.grad( loss, params, create_graphTrue, retain_graphTrue ) # 对每个参数计算对角线 diag [] for g, p in zip(grad, params): # 对梯度的每个元素求导 g_diag torch.zeros_like(p) n g.numel() for i in range(n): if g.view(-1)[i].requires_grad: second_deriv torch.autograd.grad( g.view(-1)[i], p, retain_graphTrue )[0] g_diag.view(-1)[i] second_deriv.view(-1)[i] diag.append(g_diag) return diag五、完整修复代码import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset from torch.autograd.functional import hessian as functional_hessian # Hessian 计算工具类 class HessianCalculator: Hessian 矩阵计算工具类 staticmethod def compute_full_hessian(loss, params): 计算完整的 Hessian 矩阵仅适用于小模型 复杂度: O(n²) 内存, O(n) 次反向传播 # 扁平化参数 if isinstance(params, torch.Tensor): params [params] # 计算一阶梯度 grad torch.autograd.grad( loss, params, create_graphTrue, retain_graphTrue ) flat_grad torch.cat([g.view(-1) for g in grad]) n flat_grad.numel() hessian torch.zeros(n, n) for i in range(n): grad2 torch.autograd.grad( flat_grad[i], params, retain_graphTrue ) flat_grad2 torch.cat([g.view(-1) for g in grad2]) hessian[i] flat_grad2 return hessian staticmethod def hvp(loss, params, vector): 计算 Hessian-Vector Product 复杂度: O(n) 内存, 2 次反向传播 if isinstance(params, torch.Tensor): params [params] grad torch.autograd.grad( loss, params, create_graphTrue, retain_graphTrue ) grad_v_product sum( (g * v).sum() for g, v in zip(grad, vector) ) hvp_result torch.autograd.grad( grad_v_product, params, retain_graphTrue ) return hvp_result staticmethod def compute_hessian_diagonal(model, loss_fn, inputs, targets): 计算 Hessian 对角线 outputs model(inputs) loss loss_fn(outputs, targets) params [p for p in model.parameters() if p.requires_grad] grad torch.autograd.grad( loss, params, create_graphTrue, retain_graphTrue ) diag [] for g, p in zip(grad, params): g_diag torch.zeros_like(p) n g.numel() for i in range(n): second_deriv torch.autograd.grad( g.view(-1)[i], p, retain_graphTrue, create_graphFalse )[0] g_diag.view(-1)[i] second_deriv.view(-1)[i] diag.append(g_diag) return diag staticmethod def compute_hessian_eigenvalues(hessian_matrix, top_k5): 计算 Hessian 矩阵的特征值 eigenvalues torch.linalg.eigvalsh(hessian_matrix) # 排序绝对值从大到小 sorted_vals, indices eigenvalues.abs().sort(descendingTrue) return eigenvalues[indices[:top_k]] staticmethod def lanczos_iteration(loss, params, num_iterations10): 使用 Lanczos 算法近似 Hessian 的最大/最小特征值 适用于大模型只需要 HVP 操作 n sum(p.numel() for p in params) # 初始化 v_prev torch.zeros(n) v_curr torch.randn(n) v_curr v_curr / v_curr.norm() alphas [] betas [] for i in range(num_iterations): # 重塑向量为参数形状 vector [] offset 0 for p in params: size p.numel() vector.append(v_curr[offset:offsetsize].view_as(p)) offset size # 计算 HVP hvp_result HessianCalculator.hvp(loss, params, vector) hv_flat torch.cat([h.view(-1) for h in hvp_result]) # Lanczos 递推 alpha v_curr hv_flat alphas.append(alpha.item()) hv_flat hv_flat - alpha * v_curr - (betas[-1] if betas else 0) * v_prev beta hv_flat.norm() if beta 1e-10: break betas.append(beta.item()) v_prev v_curr v_curr hv_flat / beta # 从三对角矩阵特征值近似 Hessian 特征值 T torch.diag(torch.tensor(alphas)) for i in range(len(betas)): T[i, i1] betas[i] T[i1, i] betas[i] eigenvalues torch.linalg.eigvalsh(T) return eigenvalues # 应用示例 class HessianAnalysis: Hessian 分析应用 staticmethod def analyze_loss_landscape(model, loss_fn, inputs, targets): 分析损失曲面 print( * 60) print(Hessian 分析) print( * 60) # 前向传播 outputs model(inputs) loss loss_fn(outputs, targets) params [p for p in model.parameters() if p.requires_grad] n_params sum(p.numel() for p in params) print(f参数数量: {n_params}) if n_params 100: # 小模型: 计算完整 Hessian H HessianCalculator.compute_full_hessian(loss, params) print(fHessian 形状: {H.shape}) print(fHessian 范数: {H.norm():.4f}) print(fHessian 迹: {H.trace():.4f}) # 特征值分析 eigenvalues torch.linalg.eigvalsh(H) print(f最大特征值: {eigenvalues.max():.4f}) print(f最小特征值: {eigenvalues.min():.4f}) print(f条件数: {eigenvalues.max() / eigenvalues.min():.4f}) # 正定性检查 positive (eigenvalues 0).sum().item() negative (eigenvalues 0).sum().item() zero (eigenvalues.abs() 1e-6).sum().item() print(f正特征值: {positive}, 负特征值: {negative}, 零特征值: {zero}) if negative 0: print(→ Hessian 正定凸区域) elif positive 0: print(→ Hessian 负定凹区域) else: print(→ Hessian 不定鞍点区域) else: # 大模型: 使用 HVP 和 Lanczos print(模型较大使用 HVP 近似...) # HVP 测试 vector [torch.randn_like(p) for p in params] hvp_result HessianCalculator.hvp(loss, params, vector) hvp_norm sum(h.norm() for h in hvp_result) print(fHVP 范数: {hvp_norm:.4f}) # Lanczos 近似特征值 eigenvalues HessianCalculator.lanczos_iteration( loss, params, num_iterationsmin(10, n_params) ) print(fLanczos 近似特征值: {eigenvalues}) staticmethod def optimal_brain_damage_pruning(model, loss_fn, inputs, targets, prune_ratio0.1): 使用 Hessian 对角线进行 Optimal Brain Damage 剪枝 print(\n * 60) print(Optimal Brain Damage 剪枝) print( * 60) # 计算 Hessian 对角线 diag HessianCalculator.compute_hessian_diagonal( model, loss_fn, inputs, targets ) # 计算 saliency: (weight)² × hessian_diag saliencies [] for p, h in zip(model.parameters(), diag): s (p.data ** 2) * h saliencies.append(s.view(-1)) all_saliencies torch.cat(saliencies) # 确定剪枝阈值 threshold torch.quantile( all_saliencies.abs(), prune_ratio ) # 执行剪枝 total_pruned 0 total_params 0 for p, h in zip(model.parameters(), diag): mask ((p.data ** 2) * h).abs() threshold p.data[~mask] 0 total_pruned (~mask).sum().item() total_params mask.numel() print(f剪枝比例: {100.*total_pruned/total_params:.1f}%) print(f剪枝数量: {total_pruned}/{total_params}) # 完整示例 def run_hessian_examples(): 运行 Hessian 计算示例 # 示例 1: 简单函数的 Hessian print( * 60) print(示例 1: 简单函数的 Hessian) print( * 60) def quadratic(x): return x[0]**2 3*x[0]*x[1] 2*x[1]**2 point torch.tensor([1.0, 2.0]) H functional_hessian(quadratic, point) print(f函数: f(x,y) x² 3xy 2y²) print(f在点 {point} 的 Hessian:\n{H}) print(f理论值: [[2, 3], [3, 4]]) # 示例 2: 神经网络的 Hessian print(\n * 60) print(示例 2: 神经网络的 Hessian 分析) print( * 60) # 创建小模型 model nn.Sequential( nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3) ) inputs torch.randn(16, 4) targets torch.randint(0, 3, (16,)) loss_fn nn.CrossEntropyLoss() # 分析 HessianAnalysis.analyze_loss_landscape(model, loss_fn, inputs, targets) # 示例 3: HVP print(\n * 60) print(示例 3: Hessian-Vector Product) print( * 60) # 创建较大模型 large_model nn.Sequential( nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 10) ) inputs torch.randn(32, 100) targets torch.randint(0, 10, (32,)) outputs large_model(inputs) loss nn.CrossEntropyLoss()(outputs, targets) params [p for p in large_model.parameters() if p.requires_grad] n_params sum(p.numel() for p in params) print(f模型参数量: {n_params}) # HVP vector [torch.randn_like(p) for p in params] hvp_result HessianCalculator.hvp(loss, params, vector) print(fHVP 计算成功) for i, h in enumerate(hvp_result): print(f 参数 {i}: shape{h.shape}, norm{h.norm():.4f}) if __name__ __main__: run_hessian_examples()六、常见陷阱与注意事项陷阱 1忘记create_graphTrueimport torch x torch.tensor([1.0, 2.0], requires_gradTrue) y (x ** 2).sum() # ❌ 没有 create_graphTrue一阶梯度不参与计算图 grad torch.autograd.grad(y, x)[0] try: hessian torch.autograd.grad(grad, x)[0] except RuntimeError as e: print(f错误: {e}) # ✓ 设置 create_graphTrue grad torch.autograd.grad(y, x, create_graphTrue)[0] hessian torch.autograd.grad(grad.sum(), x)[0] print(fHessian: {hessian})陷阱 2对非标量求导# ❌ autograd.grad 的 outputs 必须是标量或需要 grad_tensors x torch.tensor([1.0, 2.0], requires_gradTrue) y x ** 2 # 向量 try: grad torch.autograd.grad(y, x)[0] except RuntimeError as e: print(f错误: {e}) # ✓ 对标量求导 loss y.sum() grad torch.autograd.grad(loss, x)[0]陷阱 3内存爆炸import torch import torch.nn as nn # ❌ 对大模型计算完整 Hessian # model nn.Linear(10000, 10000) # 1 亿参数 # Hessian 是 1亿 × 1亿 矩阵 → 不可能存储 # ✓ 使用 HVP 或对角线近似 # hvp_result hvp(loss, params, vector) # O(n) 内存陷阱 4混淆 Hessian 和梯度的外积import torch x torch.tensor([1.0, 2.0], requires_gradTrue) y (x ** 3).sum() # ❌ 梯度外积不是 Hessian grad torch.autograd.grad(y, x, create_graphTrue)[0] outer_product torch.outer(grad, grad) # 这是 ∇L ⊗ ∇L不是 Hessian # ✓ 正确计算 Hessian hessian torch.zeros(2, 2) for i in range(2): hessian[i] torch.autograd.grad(grad[i], x, retain_graphTrue)[0] print(fHessian: {hessian})陷阱 5数值精度问题import torch # 二阶导数可能放大数值误差 x torch.tensor([1e-8], requires_gradTrue) y x ** 2 grad torch.autograd.grad(y, x, create_graphTrue)[0] hessian torch.autograd.grad(grad, x)[0] print(fHessian: {hessian}) # 可能不精确 # ✓ 使用 double 精度 x torch.tensor([1e-8], requires_gradTrue, dtypetorch.float64) y x ** 2 grad torch.autograd.grad(y, x, create_graphTrue)[0] hessian torch.autograd.grad(grad, x)[0] print(fHessian (double): {hessian}) # 更精确七、总结方法内存复杂度时间复杂度适用规模推荐度functional.hessianO(n²)O(n²)n 100⭐⭐⭐手动autograd.grad循环O(n²)O(n²)n 1000⭐⭐HVPO(n)O(n)任意⭐⭐⭐⭐⭐Hessian 对角线O(n)O(n²)n 10000⭐⭐⭐⭐Lanczos 迭代O(kn)O(kn)任意⭐⭐⭐⭐选择决策树需要 Hessian 信息 ├── 参数量 100 │ └── 使用 functional.hessian 或手动循环 ├── 参数量 10000 │ └── 计算 Hessian 对角线 ├── 只需要 H × v │ └── 使用 HVP └── 需要特征值 └── 使用 Lanczos 迭代核心原则对于大规模神经网络永远不要尝试计算完整的 Hessian 矩阵。使用 Hessian-Vector Product (HVP) 是处理二阶信息的标准做法它只需要 O(n) 的内存和常数次反向传播。API 速查# 一阶梯度 grad torch.autograd.grad(loss, params, create_graphTrue) # Hessian-Vector Product hvp hvp_function(loss, params, vector) # 完整 Hessian小规模 H torch.autograd.functional.hessian(func, inputs) # Hessian 对角线 diag compute_hessian_diagonal(model, loss_fn, inputs, targets) # 特征值近似 eigenvalues lanczos_iteration(loss, params, num_iterations10)通过掌握 Hessian 矩阵的计算方法可以深入理解神经网络的损失曲面结构为二阶优化、模型压缩、不确定性估计等高级应用奠定基础。
阅读完成 · 觉得有帮助？
