Custom Autograd Functions

When the built-in ops aren't enough

For 99% of model code, the built-in differentiable operations are all you need — autograd composes their derivatives via chain rule, and you never write a backward by hand. The remaining 1% is where you reach for torch.autograd.Function:

You're implementing an op whose forward isn't built from existing differentiable ops (e.g. you wrote a custom CUDA kernel).
You want to redefine the gradient (the straight-through estimator, used in quantization-aware training and discrete-output models, is the textbook example).
You want to save memory by recomputing instead of storing — though for that, torch.utils.checkpoint is usually a better fit.

The contract

A custom Function has two static methods:

forward(ctx, *inputs) — runs the forward computation. Use ctx.save_for_backward(...) to stash anything backward will need.
backward(ctx, *grad_outputs) — runs the backward computation. Returns one gradient per input to forward (or None for inputs that don't require grad).

You don't call forward or backward directly — you call MyFunction.apply(...), which sets up the graph node properly.

Code

A custom ReLU — for understanding, not production·python

import torch

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)             # stash x for backward
        return x.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        grad_in = grad_out.clone()
        grad_in[x < 0] = 0                    # zero gradient where x < 0
        return grad_in

my_relu = MyReLU.apply

x = torch.randn(5, requires_grad=True)
y = my_relu(x)
y.sum().backward()
print(x.grad)   # zeros where x<0, ones where x>=0

Straight-Through Estimator — pass gradient through a step·python

import torch

class STE(torch.autograd.Function):
    """Forward: hard threshold. Backward: pretend it was identity."""
    @staticmethod
    def forward(ctx, x):
        return (x > 0).float()

    @staticmethod
    def backward(ctx, grad_out):
        return grad_out                      # straight through

binarize = STE.apply

# Useful for quantization-aware training:
# the forward step is non-differentiable, but we still need a learning signal
x = torch.randn(4, requires_grad=True)
y = binarize(x)
y.sum().backward()
print(x.grad)   # ones — gradient passed through unchanged

Sanity-check your derivative with gradcheck·python

import torch
from torch.autograd import gradcheck

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_out):
        (x,) = ctx.saved_tensors
        grad_in = grad_out.clone()
        grad_in[x < 0] = 0
        return grad_in

# Use float64 for numerical stability when gradchecking
x = torch.randn(8, dtype=torch.float64, requires_grad=True)
print(gradcheck(MyReLU.apply, (x,), eps=1e-6, atol=1e-5))
# True if your analytic backward matches finite-difference numerics

When the built-in ops aren't enough

The contract

Code

External links

Exercise

Progress

Comments 0