How to compute gradients with TensorFlow GradientTape

Custom training code needs a way to connect a loss value back to the tensors and variables that produced it. TensorFlow uses tf.GradientTape for that automatic differentiation path, which is the piece used when debugging gradients, writing a manual training step, or checking why an optimizer update is not changing a variable.

A tape records differentiable operations while its context is active. Trainable tf.Variable objects are watched automatically when they are read inside the tape, while plain tf.Tensor inputs need tape.watch() before the calculation.

Call tape.gradient() after the forward calculation and loss are complete, then pass the returned gradients to an optimizer with the matching variables. A non-persistent tape is consumed by one gradient call, so use a fresh tape for each training step unless the project needs a persistent tape for multiple gradient queries.

1. Save the demo as

   gradient_tape_demo.py

   .

   gradient_tape_demo.py

   import os
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("ERROR")
    
   x = tf.constant(3.0)
    
   with tf.GradientTape() as tape:
       tape.watch(x)
       y = x * x + 2.0 * x
    
   dy_dx = tape.gradient(y, x)
    
   weight = tf.Variable(2.0, name="weight")
   feature = tf.constant(3.0)
   target = tf.constant(15.0)
   optimizer = tf.keras.optimizers.SGD(learning_rate=0.05)
    
   weight_before = float(weight.numpy())
    
   with tf.GradientTape() as tape:
       prediction = weight * feature
       loss = tf.square(prediction - target)
    
   gradients = tape.gradient(loss, [weight])
   loss_before = float(loss.numpy())
   optimizer.apply_gradients(zip(gradients, [weight]))
    
   prediction_after = weight * feature
   loss_after = tf.square(prediction_after - target)
    
   print(f"TensorFlow {tf.__version__}")
   print(f"dy/dx at x=3.0: {float(dy_dx.numpy()):.2f}")
   print(f"weight before: {weight_before:.2f}")
   print(f"loss before: {loss_before:.2f}")
   print(f"gradient for weight: {float(gradients[0].numpy()):.2f}")
   print(f"weight after: {float(weight.numpy()):.2f}")
   print(f"loss after: {float(loss_after.numpy()):.2f}")
   print(f"gradient connected: {gradients[0] is not None}")

   The first tape watches a plain tensor explicitly. The second tape reads a trainable variable, so TensorFlow watches that variable without a separate tape.watch() call.

2. Run the demo and confirm that the scalar derivative, variable gradient, changed weight, and lower loss all print.

   $ python3 gradient_tape_demo.py
   TensorFlow 2.21.0
   dy/dx at x=3.0: 8.00
   weight before: 2.00
   loss before: 81.00
   gradient for weight: -54.00
   weight after: 4.70
   loss after: 0.81
   gradient connected: True

3. Put the project forward pass and loss calculation inside the tape context.

   with tf.GradientTape() as tape:
       predictions = model(features, training=True)
       loss = loss_fn(labels, predictions)

   Trainable variables read by the model are watched automatically; call tape.watch(features) only when gradients are needed with respect to an input tensor.

4. Request gradients from the completed loss to the variables that should be updated.

   gradients = tape.gradient(loss, model.trainable_variables)

   None gradients mean at least one variable was not connected to the loss or was not watched by the tape. Fix that connection before applying optimizer updates.

5. Apply gradients to the same variable list.

   optimizer.apply_gradients(zip(gradients, model.trainable_variables))

   Related: How to run a custom training loop in TensorFlow

6. Recreate the tape for the next batch or use persistent=True only when multiple gradient calls must read the same recorded calculation.

   with tf.GradientTape() as tape:
       next_predictions = model(next_features, training=True)
       next_loss = loss_fn(next_labels, next_predictions)

   A non-persistent tape releases its recorded operations after tape.gradient(), so a second gradient request from the same tape raises an error.

7. Remove the demo script after the project code prints connected gradients and an optimizer update.

   $ rm gradient_tape_demo.py