GradientTape
TensorFlow records operations on a GradientTape and replays them to compute gradients. It’s the engine under every Keras optimiser; understanding it lets you write custom training loops.
Tape, gradient, optimiser step
EXAMPLE
import tensorflow as tf
W = tf.Variable(tf.random.normal((3, 1)))
b = tf.Variable(tf.zeros((1,)))
x = tf.constant([[1., 2., 3.], [4., 5., 6.]])
y = tf.constant([[10.], [20.]])
# Manual training step
optimiser = tf.keras.optimizers.Adam(0.01)
for step in range(100):
with tf.GradientTape() as tape:
pred = x @ W + b
loss = tf.reduce_mean(tf.square(pred - y))
grads = tape.gradient(loss, [W, b])
optimiser.apply_gradients(zip(grads, [W, b]))
if step % 10 == 0:
print(f'step {step:3d} loss {loss.numpy():.4f}')
# Higher-order gradients — nest tapes
with tf.GradientTape() as t2:
with tf.GradientTape() as t1:
f = W * W * W # f = W^3
g = t1.gradient(f, W) # df/dW = 3W²
h = t2.gradient(g, W) # d²f/dW² = 6W
Why it matters
Variables are watched automatically. Constants and Python ints aren’t — call tape.watch(x) if you need gradients with respect to a constant tensor.
Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.
Example
Example
import tensorflow as tf
x = tf.Variable(3.0)
with tf.GradientTape() as tape:
y = x * x + 2 * x
grad = tape.gradient(y, x)
print(grad.numpy()) # dy/dx = 2x + 2 = 8
Try it Yourself »
Exercise
Open an autodiff scope.
with tf.
() as tape:
y = x * x
PascalCase.
Discussion
Loading…