Every custom training step, and much research code, is built around the gradient tape.
TensorFlow computes gradients by recording operations. Inside a with tf.GradientTape() as tape: block, operations on watched values are written to the tape; afterwards tape.gradient(target, sources) plays them backwards to get d(target)/d(source).
What gets watched:
- Trainable
tf.Variables are watched automatically. - Constant tensors are not; call
tape.watch(t)for them.
sources can be a single variable, a list, or a nested structure (such as model.trainable_variables), and the result has the same structure. If a source did not influence the target, its gradient is None, which usually means a bug such as using a tensor copy instead of the variable, or a non-differentiable step like casting to int.
A tape is released after one gradient call. Pass persistent=True to compute several gradients from the same recording (and del tape afterwards).
Applying gradients is up to you: either update variables with assign_sub, or hand zip(grads, variables) to an optimizer's apply_gradients. When the target is not a scalar, tape.gradient returns the gradient of its sum; for full Jacobians use tape.jacobian.
Example
import tensorflow as tf
x = tf.Variable(3.0)
with tf.GradientTape() as tape:
y = x ** 2 + 2 * x
print(tape.gradient(y, x).numpy()) # 8.0 (2x + 2)
c = tf.constant(3.0)
with tf.GradientTape() as tape:
tape.watch(c) # constants need an explicit watch
y = c ** 3
print(tape.gradient(y, c).numpy()) # 27.0
w = tf.Variable([1.0, -2.0])
b = tf.Variable(0.5)
with tf.GradientTape() as tape:
loss = tf.reduce_sum(w * [3.0, 4.0]) + b
dw, db = tape.gradient(loss, [w, b])
print(dw.numpy(), db.numpy()) # [3. 4.] 1.0
w.assign_sub(0.1 * dw) # one gradient-descent step
print(w.numpy()) # [ 0.7 -2.4]Remember
- Record inside
tf.GradientTape, thentape.gradient(target, sources) - Variables are watched automatically; constants need
tape.watch - A
Nonegradient means the source did not affect the target
Watch out for
- Calling
tape.gradienttwice on a non-persistent tape, which raises an error.