Tensorflow 2: GradientTape returns None

Viewed 99

I have the following code, written in tf.keras with Tensorflow 2. Basically; I need the cross entropy term's gradient with respect to the variable self.temperature. dce1_dx correctly calculates the derivative. But on the other hand, when I wrap the same cross entropy calculation into a tf.keras.Model object, the second gradient calculation, dce2_dx returns None. What is the difference between these two tf.GradientTape calculations? I am experienced in TF1 but newly switching to TF2 and eager execution, so I am stuck at that point.

import numpy as np
import tensorflow as tf

logits = np.random.uniform(low=-10.0, high=10.0, size=(10000, 5))
labels = np.random.randint(low=0, high=5, size=(10000, ))
logits_tf = tf.keras.Input(name="logits_tf", shape=(logits.shape[1]), dtype=tf.float32)
labels_tf = tf.keras.Input(name="labels_tf", shape=(), dtype=tf.int32)
dataset = tf.data.Dataset.from_tensor_slices((logits, labels))
dataset = dataset.batch(batch_size=logits.shape[0])
for lgts, idx in dataset:
   temperature = tf.Variable(name="temperature", dtype=tf.float32, initial_value=tf.constant(2.0),
                               trainable=True)
   scaled_logits = logits_tf / temperature
   ce_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
   ce_loss = ce_loss(labels_tf, scaled_logits)
   model = tf.keras.Model(inputs=[logits_tf, labels_tf], outputs=[ce_loss], name="calibration_model")
   with tf.GradientTape() as tape0:
       tape0.watch(temperature)
       scaled_lgts = tf.cast(lgts, tf.float32) / temperature
       ce = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
       ce = ce(idx, scaled_lgts)
   dce1_dx = tape0.gradient(ce, temperature)

   with tf.GradientTape() as tape1:
       # Compute the derivative: d{CrossEntropy}/d{Temperature}
       tape1.watch(temperature)
       ce2 = model([lgts, idx])
   # !!!Returns None!!!
   dce2_dx = tape1.gradient(ce2, temperature)
0 Answers
Related