【Bug已解决】How to use autograd.gradcheck in PyTorch? 解决方案问题描述torch.autograd.gradcheck是 PyTorch 中用于验证自定义 autograd Function 梯度正确性的工具。它通过数值微分numerical differentiation计算参考梯度然后与 autograd 计算的解析梯度进行比较。许多开发者在使用gradcheck时会遇到以下问题要求float64精度gradcheck不接受float32输入输入需要requires_gradTrue忘记设置导致报错多输入/多输出函数的用法如何传入多个张量数值不稳定某些函数在数值微分时不稳定gradgradcheck的使用验证二阶梯度的正确性典型错误import torch # 错误使用 float32 x torch.randn(5, requires_gradTrue) # 默认 float32 # torch.autograd.gradcheck(torch.exp, x) # RuntimeError: gradcheck expects all tensors to be of dtype torch.float64 # 错误忘记 requires_grad x torch.randn(5, dtypetorch.float64) # 没有 requires_grad # torch.autograd.gradcheck(torch.exp, x) # RuntimeError: gradcheck requires input tensors to have requires_gradTrue错误复现错误场景一float32 输入import torch x torch.randn(5, requires_gradTrue, dtypetorch.float32) # torch.autograd.gradcheck(lambda t: t.pow(2).sum(), x) # RuntimeError: gradcheck expects all tensors to be of dtype torch.float64错误场景二忘记 requires_gradx torch.randn(5, dtypetorch.float64) # 没有 requires_gradTrue # torch.autograd.gradcheck(lambda t: t.exp().sum(), x) # RuntimeError: gradcheck requires input tensors to have requires_gradTrue错误场景三多输入函数x torch.randn(5, dtypetorch.float64, requires_gradTrue) y torch.randn(5, dtypetorch.float64, requires_gradTrue) # 错误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, dtypetorch.float64, requires_gradTrue) # 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 eps1e-6, # 有限差分的步长 atol1e-5, # 绝对容差 rtol1e-3, # 相对容差 raise_exceptionTrue, # 不匹配时是否抛出异常 )4.gradgradcheck的原理gradgradcheck验证二阶梯度的正确性。它对gradcheck的结果再应用一次gradcheck即验证 backward 函数本身的梯度。解决方案方案一基本用法import torch # 正确的 gradcheck 用法 x torch.randn(5, dtypetorch.float64, requires_gradTrue) # 检查 torch.exp 的梯度 result torch.autograd.gradcheck(torch.exp, (x,)) print(ftorch.exp gradcheck: {result}) # True # 检查自定义函数 def my_square(x): return x.pow(2).sum() result torch.autograd.gradcheck(my_square, (x,)) print(fmy_square gradcheck: {result}) # True方案二多输入函数import torch x torch.randn(5, dtypetorch.float64, requires_gradTrue) y torch.randn(5, dtypetorch.float64, requires_gradTrue) # 多输入函数输入必须是 tuple def my_func(a, b): return (a * b a.exp() b.log()).sum() result torch.autograd.gradcheck(my_func, (x, y)) print(fMulti-input gradcheck: {result}) # True方案三自定义 autograd Functionimport torch class MyReLU(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x.clamp(min0) staticmethod def backward(ctx, grad_output): x, ctx.saved_tensors grad_x grad_output * (x 0).float() return grad_x # 测试 x torch.randn(10, dtypetorch.float64, requires_gradTrue) result torch.autograd.gradcheck(lambda t: MyReLU.apply(t).sum(), (x,)) print(fMyReLU 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, dtypetorch.float64, requires_gradTrue) # 测试各种函数 functions [ (exp, lambda t: t.exp().sum()), (log, lambda t: t.clamp(min0.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_slope0.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(dimdim, keepdimTrue)[0] exp_x torch.exp(x - x_max) sum_exp exp_x.sum(dimdim, keepdimTrue) 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(dimctx.dim, keepdimTrue) 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, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float64, requires_gradTrue) b torch.randn(4, 5, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float64, requires_gradTrue) y torch.randn(5, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float64, requires_gradTrue) b torch.randn(4, 2, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float32, requires_gradTrue) try: torch.autograd.gradcheck(lambda t: t.exp().sum(), (x_f32,)) except RuntimeError as e: print(f 错误: {str(e)[:70]}) print(f 修复: 使用 dtypetorch.float64) # 错误 2: 没有 requires_grad print(\n2. 缺少 requires_grad:) x_no_grad torch.randn(5, dtypetorch.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_gradTrue) # 错误 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, dtypetorch.float64, requires_gradTrue) 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(nameNone): 创建 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, gamma2.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(min1e-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, dtypetorch.float64, requires_gradTrue) target torch.randint(0, 2, (10,), dtypetorch.float64) # 使用 PyTorch 内置操作实现让 autograd 自动推导 def focal_loss_autograd(logits, target, gamma2.0, alpha0.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(min1e-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, dtypetorch.float64, requires_gradTrue) k torch.randn(2, 3, 8, dtypetorch.float64, requires_gradTrue) v torch.randn(2, 3, 8, dtypetorch.float64, requires_gradTrue) 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, dtypetorch.float32, requires_gradTrue) # 正确 x torch.randn(5, dtypetorch.float64, requires_gradTrue)2. 输入必须是 tuple# 错误 torch.autograd.gradcheck(func, x) # 正确 torch.autograd.gradcheck(func, (x,))3. 函数必须返回标量或张量# gradcheck 要求函数输出是张量 # 如果输出是标量gradcheck 会自动处理4.eps的选择# 默认 eps1e-6 适用于大多数情况 # 如果函数在数值上不稳定可以调整 torch.autograd.gradcheck(func, inputs, eps1e-5) # 更大的步长5. 不可微点# ReLU 在 0 处不可微 x torch.tensor([0.0], dtypetorch.float64, requires_gradTrue) # gradcheck 可能失败因为数值微分在不可微点不准确 # 避免在不可微点附近测试6.check_undefined_grad# 检查未定义梯度 torch.autograd.gradcheck(func, inputs, check_undefined_gradFalse)7.check_batched_grad# PyTorch 1.9 支持批量梯度检查 torch.autograd.gradcheck(func, inputs, check_batched_gradTrue)8.raise_exceptionFalse获取结果result torch.autograd.gradcheck(func, inputs, raise_exceptionFalse) # result 是 boolTrue 表示通过总结torch.autograd.gradcheck是验证自定义梯度实现正确性的重要工具。使用要点必须使用float64数值微分需要高精度输入需要requires_gradTrue否则无法计算梯度输入必须是 tuple(x,)而非x多输入函数gradcheck(func, (x, y, z))自定义 Function验证forward和backward的一致性gradgradcheck验证二阶梯度的正确性避免不可微点如 ReLU 在 0 处调整eps数值不稳定时增大步长在实现自定义 autograd Function 时始终使用gradcheck验证梯度的正确性可以避免训练中的隐蔽 bug确保模型正确学习。
