自己用 PyTorch 和 NumPy 结合起来写了一个扩展网络,之前都是假定一些数据手动计算来校验前向传播和反向传播的对错。最近才想起来用 PyTorch
的 gradcheck 功能。下面简单记录一下。
torch.autograd.gradcheck(func, inputs, eps=1e-06, atol=1e-05,
rtol=0.001, raise_exception=True, check_sparse_nnz=False, nondet_tol=0.0)
该函数应用小有限差分来校验相关 Tensor 的梯度,Tensor 是浮点类型且设置 requires_grad=True
。func
就是自己写的扩展网络函数;inputs
是该网络的输入;eps
是有限差分的扰动;atol
是绝对容差;rtol
是相对容差。
且需要注意,输入 inputs
需设置成 double
精度。如果不是的话,梯度校验可能会失败。
下面是使用例子,其中 HitMissLayer
是我自己写的扩展网络:
import torch
from torch.autograd.gradcheck import gradcheck
hmt = HitMissLayer([1, 1, 32, 32], 2)
hmt = hmt.double()
inputss = torch.randn((1, 1, 32, 32), dtype=torch.double, requires_grad=True)
test = gradcheck(hmt, inputss, eps=1e-6, atol=1e-4)
print("Are the gradients correct: ", test)
如果校验成功,上面代码会出现如下输出:
Are the gradients correct: True