checkpoint(function: Callable[..., Tensor], args: Tensor = (), preserve_rng_state: bool = True, use_reentrant: bool = True, kwargs: object = {})Run function under gradient checkpointing.
Executes function(*args, **kwargs) during the forward pass
without tracking intermediate activations (no_grad context).
During the backward pass the function is re-executed under
enable_grad to reconstruct the local autograd graph, and the
gradients are computed through that graph.
Parameters
functioncallableThe differentiable segment to checkpoint. Must accept tensors as
positional arguments and return a single
lucid.Tensor.Positional tensor inputs to function.
preserve_rng_statebool= TrueAccepted for API compatibility. RNG state restoration is not yet
implemented — set to
False when function contains stochastic
layers.use_reentrantbool= TrueAccepted for API compatibility. Only the reentrant (default)
implementation is provided.
**kwargsobject= {}Extra keyword arguments forwarded to function on both the
forward and recomputation passes.
Returns
TensorOutput of function(*args, **kwargs).
Examples
>>> def block(x):
... return lucid.nn.functional.relu(x @ W + b)
>>> y = lucid.autograd.checkpoint(block, x)
>>> y.sum().backward()