【Bug已解决】How to use autograd.gradcheck in PyTorch? 解决方案
问题描述
torch.autograd.gradcheck是 PyTorch 中用于验证自定义 autograd Function 梯度正确性的工具。它通过数值微分（numerical differentiation）计算参考梯度，然后与 autograd 计算的解析梯度进行比较。
许多开发者在使用gradcheck时会遇到以下问题：
- 要求
float64精度：gradcheck不接受float32输入 - 输入需要
requires_grad=True：忘记设置导致报错 - 多输入/多输出函数的用法：如何传入多个张量
- 数值不稳定：某些函数在数值微分时不稳定
gradgradcheck的使用：验证二阶梯度的正确性
典型错误：
import torch # 错误：使用 float32 x = torch.randn(5, requires_grad=True) # 默认 float32 # torch.autograd.gradcheck(torch.exp, x) # RuntimeError: gradcheck expects all tensors to be of dtype torch.float64 # 错误：忘记 requires_grad x = torch.randn(5, dtype=torch.float64) # 没有 requires_grad # torch.autograd.gradcheck(torch.exp, x) # RuntimeError: gradcheck requires input tensors to have requires_grad=True错误复现
错误场景一：float32 输入
import torch x = torch.randn(5, requires_grad=True, dtype=torch.float32) # torch.autograd.gradcheck(lambda t: t.pow(2).sum(), x) # RuntimeError: gradcheck expects all tensors to be of dtype torch.float64错误场景二：忘记 requires_grad
x = torch.randn(5, dtype=torch.float64) # 没有 requires_grad=True # torch.autograd.gradcheck(lambda t: t.exp().sum(), x) # RuntimeError: gradcheck requires input tensors to have requires_grad=True错误场景三：多输入函数
x = torch.randn(5, dtype=torch.float64, requires_grad=True) y = torch.randn(5, dtype=torch.float64, requires_grad=True) # 错误：gradcheck 需要输入是 tuple # torch.autograd.gradcheck(lambda a, b: (a * b).sum(), x, y) # TypeError: gradcheck expects inputs as a tuple # 正确 # torch.autograd.gradcheck(lambda a, b: (a * b).sum(), (x, y))错误场景四：自定义 Function 梯度错误
import torch class WrongSquare(torch.autograd.Function): @staticmethod def forward(ctx, x): return x ** 2 @staticmethod def backward(ctx, grad_output): # 错误：梯度应该是 2*x，但写成了 x return grad_output * x # 错误！ x = torch.randn(5, dtype=torch.float64, requires_grad=True) # torch.autograd.gradcheck(lambda t: WrongSquare.apply(t).sum(), x) # RuntimeError: Jacobian mismatch for output 0 with respect to input 0根因分析
1.gradcheck的工作原理
gradcheck使用有限差分法（finite differences）计算数值梯度：
数值梯度 ≈ (f(x + ε) - f(x - ε)) / (2ε)其中ε是一个很小的值（默认1e-6）。
然后比较数值梯度和 autograd 计算的解析梯度，如果差异超过容差（atol/rtol），则报错。
2. 为什么需要float64
有限差分法涉及两个非常接近的数相减，会产生严重的精度损失（catastrophic cancellation）。float32只有约 7 位有效数字，无法提供足够的精度来计算可靠的数值梯度。float64有约 15-17 位有效数字，可以满足要求。
3.gradcheck的参数
torch.autograd.gradcheck( fn, # 要检查的函数 inputs, # 输入张量的 tuple eps=1e-6, # 有限差分的步长 atol=1e-5, # 绝对容差 rtol=1e-3, # 相对容差 raise_exception=True, # 不匹配时是否抛出异常 )4.gradgradcheck的原理
gradgradcheck验证二阶梯度的正确性。它对gradcheck的结果再应用一次gradcheck，即验证 backward 函数本身的梯度。
解决方案
方案一：基本用法
import torch # 正确的 gradcheck 用法 x = torch.randn(5, dtype=torch.float64, requires_grad=True) # 检查 torch.exp 的梯度 result = torch.autograd.gradcheck(torch.exp, (x,)) print(f"torch.exp gradcheck: {result}") # True # 检查自定义函数 def my_square(x): return x.pow(2).sum() result = torch.autograd.gradcheck(my_square, (x,)) print(f"my_square gradcheck: {result}") # True方案二：多输入函数
import torch x = torch.randn(5, dtype=torch.float64, requires_grad=True) y = torch.randn(5, dtype=torch.float64, requires_grad=True) # 多输入函数：输入必须是 tuple def my_func(a, b): return (a * b + a.exp() + b.log()).sum() result = torch.autograd.gradcheck(my_func, (x, y)) print(f"Multi-input gradcheck: {result}") # True方案三：自定义 autograd Function
import torch class MyReLU(torch.autograd.Function): @staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x.clamp(min=0) @staticmethod def backward(ctx, grad_output): x, = ctx.saved_tensors grad_x = grad_output * (x > 0).float() return grad_x # 测试 x = torch.randn(10, dtype=torch.float64, requires_grad=True) result = torch.autograd.gradcheck(lambda t: MyReLU.apply(t).sum(), (x,)) print(f"MyReLU gradcheck: {result}") # True完整修复代码
""" PyTorch autograd.gradcheck 完整使用示例 涵盖基本函数、自定义 Function、多输入/输出、二阶梯度等 """ import torch import torch.nn as nn import torch.nn.functional as F import math # ============================================================ # 第一部分：基本 gradcheck 用法 # ============================================================ def demo_basic_gradcheck(): """基本 gradcheck 用法""" print("=" * 60) print("基本 gradcheck") print("=" * 60) # 创建 float64 输入 x = torch.randn(5, dtype=torch.float64, requires_grad=True) # 测试各种函数 functions = [ ("exp", lambda t: t.exp().sum()), ("log", lambda t: t.clamp(min=0.1).log().sum()), # 避免 log(0) ("square", lambda t: t.pow(2).sum()), ("cube", lambda t: t.pow(3).sum()), ("sigmoid", lambda t: torch.sigmoid(t).sum()), ("tanh", lambda t: torch.tanh(t).sum()), ("relu", lambda t: F.relu(t).sum()), ] for name, func in functions: try: result = torch.autograd.gradcheck(func, (x,)) print(f" {name:15s}: {'PASS' if result else 'FAIL'}") except RuntimeError as e: print(f" {name:15s}: ERROR - {str(e)[:60]}") # ============================================================ # 第二部分：自定义 autograd Function 验证 # ============================================================ class LeakyReLUManual(torch.autograd.Function): """手动实现的 LeakyReLU""" @staticmethod def forward(ctx, x, negative_slope=0.01): ctx.save_for_backward(x) ctx.negative_slope = negative_slope return torch.where(x >= 0, x, x * negative_slope) @staticmethod def backward(ctx, grad_output): x, = ctx.saved_tensors grad_x = torch.where(x >= 0, grad_output, grad_output * ctx.negative_slope) return grad_x, None # None 对应 negative_slope 的梯度 class SoftmaxManual(torch.autograd.Function): """手动实现的 Softmax""" @staticmethod def forward(ctx, x, dim=-1): x_max = x.max(dim=dim, keepdim=True)[0] exp_x = torch.exp(x - x_max) sum_exp = exp_x.sum(dim=dim, keepdim=True) output = exp_x / sum_exp ctx.save_for_backward(output) ctx.dim = dim return output @staticmethod def backward(ctx, grad_output): output, = ctx.saved_tensors # softmax 的 Jacobian: diag(s) - s * s^T # 对于向量: grad_input = s * (grad_output - sum(grad_output * s)) grad_sum = (grad_output * output).sum(dim=ctx.dim, keepdim=True) grad_input = output * (grad_output - grad_sum) return grad_input, None class MatMulManual(torch.autograd.Function): """手动实现的矩阵乘法""" @staticmethod def forward(ctx, a, b): ctx.save_for_backward(a, b) return a @ b @staticmethod def backward(ctx, grad_output): a, b = ctx.saved_tensors grad_a = grad_output @ b.t() grad_b = a.t() @ grad_output return grad_a, grad_b def demo_custom_functions(): """自定义 Function 验证""" print("\n" + "=" * 60) print("自定义 Function gradcheck") print("=" * 60) # LeakyReLU x = torch.randn(10, dtype=torch.float64, requires_grad=True) result = torch.autograd.gradcheck( lambda t: LeakyReLUManual.apply(t, 0.01).sum(), (x,) ) print(f" LeakyReLU: {'PASS' if result else 'FAIL'}") # Softmax x = torch.randn(5, 3, dtype=torch.float64, requires_grad=True) result = torch.autograd.gradcheck( lambda t: SoftmaxManual.apply(t, dim=-1).sum(), (x,) ) print(f" Softmax: {'PASS' if result else 'FAIL'}") # MatMul a = torch.randn(3, 4, dtype=torch.float64, requires_grad=True) b = torch.randn(4, 5, dtype=torch.float64, requires_grad=True) result = torch.autograd.gradcheck( lambda ta, tb: MatMulManual.apply(ta, tb).sum(), (a, b) ) print(f" MatMul: {'PASS' if result else 'FAIL'}") # ============================================================ # 第三部分：多输入/多输出函数 # ============================================================ def demo_multi_io(): """多输入/多输出 gradcheck""" print("\n" + "=" * 60) print("多输入/多输出 gradcheck") print("=" * 60) # 多输入 x = torch.randn(5, dtype=torch.float64, requires_grad=True) y = torch.randn(5, dtype=torch.float64, requires_grad=True) def multi_input(a, b): return (a * b + a.exp() - b.pow(2)).sum() result = torch.autograd.gradcheck(multi_input, (x, y)) print(f" 多输入: {'PASS' if result else 'FAIL'}") # 多输出（gradcheck 会检查所有输出对所有输入的 Jacobian） def multi_output(a, b): return (a * b).sum(), (a + b).sum(), (a - b).sum() result = torch.autograd.gradcheck(multi_output, (x, y)) print(f" 多输出: {'PASS' if result else 'FAIL'}") # 矩阵输入 a = torch.randn(3, 4, dtype=torch.float64, requires_grad=True) b = torch.randn(4, 2, dtype=torch.float64, requires_grad=True) def matmul_func(x, w): return (x @ w).sum() result = torch.autograd.gradcheck(matmul_func, (a, b)) print(f" 矩阵乘法: {'PASS' if result else 'FAIL'}") # ============================================================ # 第四部分：gradgradcheck（二阶梯度验证） # ============================================================ def demo_gradgradcheck(): """二阶梯度验证""" print("\n" + "=" * 60) print("gradgradcheck（二阶梯度）") print("=" * 60) # 简单函数 x = torch.randn(5, dtype=torch.float64, requires_grad=True) def my_func(t): return t.pow(3).sum() # f(x) = x^3, f'(x) = 3x^2, f''(x) = 6x result = torch.autograd.gradgradcheck(my_func, (x,)) print(f" x^3 二阶梯度: {'PASS' if result else 'FAIL'}") # 自定义 Function 的二阶梯度 class CubicFunction(torch.autograd.Function): @staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x.pow(3) @staticmethod def backward(ctx, grad_output): x, = ctx.saved_tensors return grad_output * 3 * x.pow(2) # 导数 3x^2 x = torch.randn(5, dtype=torch.float64, requires_grad=True) result = torch.autograd.gradgradcheck( lambda t: CubicFunction.apply(t).sum(), (x,) ) print(f" 自定义 Cubic 二阶: {'PASS' if result else 'FAIL'}") # ============================================================ # 第五部分：常见错误与调试 # ============================================================ def demo_common_errors(): """常见错误演示""" print("\n" + "=" * 60) print("常见错误与修复") print("=" * 60) # 错误 1: float32 print("1. float32 输入:") x_f32 = torch.randn(5, dtype=torch.float32, requires_grad=True) try: torch.autograd.gradcheck(lambda t: t.exp().sum(), (x_f32,)) except RuntimeError as e: print(f" 错误: {str(e)[:70]}") print(f" 修复: 使用 dtype=torch.float64") # 错误 2: 没有 requires_grad print("\n2. 缺少 requires_grad:") x_no_grad = torch.randn(5, dtype=torch.float64) try: torch.autograd.gradcheck(lambda t: t.exp().sum(), (x_no_grad,)) except RuntimeError as e: print(f" 错误: {str(e)[:70]}") print(f" 修复: 设置 requires_grad=True") # 错误 3: 错误的梯度实现 print("\n3. 错误的梯度实现:") class WrongGrad(torch.autograd.Function): @staticmethod def forward(ctx, x): return x.pow(2) @staticmethod def backward(ctx, grad_output): return grad_output * x # 错误！应该是 2*x x = torch.randn(5, dtype=torch.float64, requires_grad=True) try: # 需要保存 x 到 ctx 才能在 backward 中使用 class WrongGradFixed(torch.autograd.Function): @staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x.pow(2) @staticmethod def backward(ctx, grad_output): x, = ctx.saved_tensors return grad_output * x # 故意写错：应该是 2*x torch.autograd.gradcheck(lambda t: WrongGradFixed.apply(t).sum(), (x,)) except RuntimeError as e: print(f" 错误: Jacobian mismatch") print(f" 修复: backward 应返回 grad_output * 2 * x") # ============================================================ # 第六部分：实用工具 # ============================================================ def create_gradcheck_wrapper(name=None): """创建 gradcheck 包装器""" def gradcheck_decorator(func): """装饰器：自动将输入转为 float64 并运行 gradcheck""" def wrapper(*args, **kwargs): # 转换输入为 float64 test_args = [] for arg in args: if isinstance(arg, torch.Tensor): test_arg = arg.detach().double().requires_grad_(True) test_args.append(test_arg) else: test_args.append(arg) # 运行 gradcheck try: result = torch.autograd.gradcheck( lambda *a: func(*a, **kwargs) if isinstance(func(*a, **kwargs), torch.Tensor) else func(*a, **kwargs), tuple(test_args) ) label = name or func.__name__ print(f" [gradcheck] {label}: {'PASS' if result else 'FAIL'}") return func(*args, **kwargs) except Exception as e: print(f" [gradcheck] FAILED: {e}") raise return wrapper return gradcheck_decorator def demo_practical_usage(): """实际使用示例""" print("\n" + "=" * 60) print("实际使用示例") print("=" * 60) # 验证自定义损失函数 class FocalLossManual(torch.autograd.Function): """手动实现的 Focal Loss""" @staticmethod def forward(ctx, x, target, gamma=2.0): # x: logits, target: binary labels p = torch.sigmoid(x) pt = p * target + (1 - p) * (1 - target) alpha = 0.25 at = alpha * target + (1 - alpha) * (1 - target) loss = -at * (1 - pt).pow(gamma) * torch.log(pt.clamp(min=1e-8)) ctx.save_for_backward(p, target, pt, at) ctx.gamma = gamma return loss.sum() @staticmethod def backward(ctx, grad_output): p, target, pt, at = ctx.saved_tensors gamma = ctx.gamma # 数值梯度验证比较复杂，这里简化 # 实际使用时建议用 autograd 自动推导 return grad_output * torch.ones_like(p), None, None # 验证 x = torch.randn(10, dtype=torch.float64, requires_grad=True) target = torch.randint(0, 2, (10,), dtype=torch.float64) # 使用 PyTorch 内置操作实现，让 autograd 自动推导 def focal_loss_autograd(logits, target, gamma=2.0, alpha=0.25): p = torch.sigmoid(logits) pt = p * target + (1 - p) * (1 - target) at = alpha * target + (1 - alpha) * (1 - target) loss = -at * (1 - pt).pow(gamma) * torch.log(pt.clamp(min=1e-8)) return loss.sum() result = torch.autograd.gradcheck( lambda t: focal_loss_autograd(t, target), (x,) ) print(f" Focal Loss (autograd): {'PASS' if result else 'FAIL'}") # 验证注意力机制 def attention(q, k, v): """Scaled dot-product attention""" d_k = k.size(-1) scores = q @ k.transpose(-2, -1) / math.sqrt(d_k) attn = F.softmax(scores, dim=-1) return (attn @ v).sum() q = torch.randn(2, 3, 8, dtype=torch.float64, requires_grad=True) k = torch.randn(2, 3, 8, dtype=torch.float64, requires_grad=True) v = torch.randn(2, 3, 8, dtype=torch.float64, requires_grad=True) result = torch.autograd.gradcheck(attention, (q, k, v)) print(f" Attention: {'PASS' if result else 'FAIL'}") # ============================================================ # 主函数 # ============================================================ if __name__ == '__main__': demo_basic_gradcheck() demo_custom_functions() demo_multi_io() demo_gradgradcheck() demo_common_errors() demo_practical_usage() print("\n" + "=" * 60) print("所有示例运行完成！") print("=" * 60)常见陷阱与注意事项
1. 必须使用float64
# 错误 x = torch.randn(5, dtype=torch.float32, requires_grad=True) # 正确 x = torch.randn(5, dtype=torch.float64, requires_grad=True)2. 输入必须是 tuple
# 错误 torch.autograd.gradcheck(func, x) # 正确 torch.autograd.gradcheck(func, (x,))3. 函数必须返回标量或张量
# gradcheck 要求函数输出是张量 # 如果输出是标量，gradcheck 会自动处理4.eps的选择
# 默认 eps=1e-6 适用于大多数情况 # 如果函数在数值上不稳定，可以调整 torch.autograd.gradcheck(func, inputs, eps=1e-5) # 更大的步长5. 不可微点
# ReLU 在 0 处不可微 x = torch.tensor([0.0], dtype=torch.float64, requires_grad=True) # gradcheck 可能失败，因为数值微分在不可微点不准确 # 避免在不可微点附近测试6.check_undefined_grad
# 检查未定义梯度 torch.autograd.gradcheck(func, inputs, check_undefined_grad=False)7.check_batched_grad
# PyTorch >= 1.9 支持批量梯度检查 torch.autograd.gradcheck(func, inputs, check_batched_grad=True)8.raise_exception=False获取结果
result = torch.autograd.gradcheck(func, inputs, raise_exception=False) # result 是 bool，True 表示通过总结
torch.autograd.gradcheck是验证自定义梯度实现正确性的重要工具。使用要点：
- 必须使用
float64：数值微分需要高精度 - 输入需要
requires_grad=True：否则无法计算梯度 - 输入必须是 tuple：
(x,)而非x - 多输入函数：
gradcheck(func, (x, y, z)) - 自定义 Function：验证
forward和backward的一致性 gradgradcheck：验证二阶梯度的正确性- 避免不可微点：如 ReLU 在 0 处
- 调整
eps：数值不稳定时增大步长
在实现自定义 autograd Function 时，始终使用gradcheck验证梯度的正确性，可以避免训练中的隐蔽 bug，确保模型正确学习。
