Gradients with GradientTape

TensorFlow · Essentials · 12 min

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, then tape.gradient(target, sources)
  • Variables are watched automatically; constants need tape.watch
  • A None gradient means the source did not affect the target

Watch out for

  • Calling tape.gradient twice on a non-persistent tape, which raises an error.

Check yourself

  1. tape.gradient(loss, v) returns None. What is the most likely cause?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: www.tensorflow.org/guide/autodiff