TopolearnNeeded for numerically stable functions, straight-through estimators, gradient penalties and second-order methods.
`@tf.custom_gradient` lets you define a function together with its gradient. The decorated function returns (output, grad_fn), where grad_fn(upstream) returns the gradient with respect to each input. Uses:
- Numerical stability: the textbook example is
log(1 + exp(x)), whose automatic gradient isexp(x) / (1 + exp(x))and becomesinf / inf = nanfor large x. A hand-written gradient1 - 1 / (1 + exp(x))stays finite. - Straight-through estimators and gradient clipping inside a layer.
- Saving memory by recomputing instead of storing (see also
tf.recompute_grad).
tf.stop_gradient(x) passes values through but blocks gradients, used for target networks and some losses.
Higher-order gradients come from nesting tapes: compute a gradient inside an outer tape, then differentiate it again with the outer tape. That gives second derivatives, Hessian-vector products and gradient penalties (as in WGAN-GP).
For full derivative matrices, tape.jacobian(y, x) returns the Jacobian, and tape.batch_jacobian(y, x) computes per-example Jacobians efficiently when examples are independent. Both use pfor vectorisation under the hood; pass experimental_use_pfor=False to fall back to a loop if an op is not supported.
Example
import tensorflow as tf
@tf.custom_gradient
def log1pexp(x):
e = tf.exp(x)
def grad(upstream):
return upstream * (1 - 1 / (1 + e)) # stable even when e is inf
return tf.math.log(1 + e), grad
x = tf.constant(100.0)
with tf.GradientTape() as tape:
tape.watch(x)
y = log1pexp(x)
print(tape.gradient(y, x).numpy()) # 1.0 (automatic gradient gives nan)
v = tf.Variable(3.0)
with tf.GradientTape() as outer:
with tf.GradientTape() as inner:
y = v ** 3
dy = inner.gradient(y, v) # 3v^2 = 27
d2y = outer.gradient(dy, v) # 6v = 18
print(dy.numpy(), d2y.numpy()) # 27.0 18.0
w = tf.Variable([1.0, 2.0])
with tf.GradientTape() as tape:
out = w ** 2
print(tape.jacobian(out, w).numpy()) # [[2. 0.] [0. 4.]]Remember
@tf.custom_gradientreturns (output, grad_fn)- Nest tapes for second derivatives;
jacobian/batch_jacobianfor matrices tf.stop_gradientblocks gradient flow
Watch out for
- Computing the inner gradient outside the outer tape's context, so the outer tape cannot differentiate it.