🎓 A1Learner
🔍

PyTorch Gradcheck

gradcheck is a test that checks a backward rule by nudging its inputs. The comparison runs in float64, the double-precision dtype with extra digits.

A gradient measured by nudging

A numerical gradient nudges an input a little and measures the slope from the outputs. For x squared at two, the measured slope matches the exact slope of four.

import torch
from torch.autograd import gradcheck


def f(v):
    return v * v


eps = 0.001
numeric = (f(2.0 + eps) - f(2.0 - eps)) / (2 * eps)
print("Numerical gradient:", round(numeric, 3))
print("Exact gradient:", 4.0)

A correct backward passes

The custom square runs here in float64, as gradcheck demands. gradcheck returns True, so the backward rule agrees with the nudges.

class SquareFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * x

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * 2 * x


good = torch.tensor([0.5, 1.5], dtype=torch.float64, requires_grad=True)
print("Correct backward passes:", gradcheck(SquareFn.apply, (good,)))

Float32 noise fails a correct rule

The same correct square now receives a float32 input. gradcheck raises GradcheckError, because the tiny nudges drown in float32 rounding noise.

bad32 = torch.tensor([0.5], dtype=torch.float32, requires_grad=True)
print(gradcheck(SquareFn.apply, (bad32,)))

Example

import torch
from torch.autograd import gradcheck


class SquareFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * x

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * 2 * x


class WrongSquareFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * x

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * 3 * x


good_input = torch.tensor([0.5, 1.5], dtype=torch.float64,
                          requires_grad=True)
try:
    good_ok = gradcheck(SquareFn.apply, (good_input,), raise_exception=True)
except Exception:
    good_ok = False
print("Correct backward passes:", good_ok)

bad_input = torch.tensor([0.5, 1.5], dtype=torch.float64,
                         requires_grad=True)
try:
    gradcheck(WrongSquareFn.apply, (bad_input,), raise_exception=True)
    bad_ok = True
except Exception as error:
    bad_ok = False
    print("Wrong backward error type:", type(error).__name__)
print("Wrong backward passes:", bad_ok)

Key points

  • Numerical gradient: a slope measured by nudging an input and watching the output
  • Analytic gradient: the gradient a backward rule claims is correct
  • Use torch.autograd.gradcheck with a function and a tuple of its inputs
  • gradcheck needs float64 inputs, because float32 rounding swamps the tiny nudges
  • gradcheck returns True on a pass and raises GradcheckError on a mismatch
  • In short, gradcheck trusts no backward rule until the nudges agree with it

In this example, two custom squares share one forward and differ only in backward. SquareFn multiplies the incoming gradient by two times x, which is correct. WrongSquareFn multiplies by three times x instead, which is wrong. Both inputs hold 0.5 and 1.5 in float64, as gradcheck requires. The first gradcheck call returns True, and the program prints that pass. The second call raises GradcheckError, so the program catches it in an except block. The block prints the error type name and then prints the failure as False. Catching the error lets one program report both outcomes and still exit cleanly.

Note

gradcheck suits small inputs, because every input is nudged one entry at a time

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does gradcheck do when a backward rule is wrong?