How can I fix this problem of automated differentiation with Tensorflow?

Question

I need to compute the gradient wrt a loss for the variables of a defined neural network, the loss is computed correctly but the gradients are None. The code is the following:

variables = self.model.trainable_variables
for var in variables:
   print("{}".format(var.name)
with tf.GradientTape() as tape:

   y = self.model.predict(states[t])
   y = tf.reshape(y, constants.NUM_ACTIONS)
   loss = y[actions[t]]

grads = tape.gradient(loss, variables)

print("Information about gradients")
for var, g in zip(variables, grads):
   print(f'{var.name}, shape: {g.shape}')

I get the following error:

AttributeError: 'NoneType' object has no attribute 'shape'

How can I fix this problem?

Answer 1

Direct indexing y[actions[t]] is not differentiable, so your loss tensor is disconnected from the graph.

You can rewrite your code to produce the loss for all elements, but use an input binary mask to select only one element. This way your final loss is connected to all inputs (states[t] and mask[t]).

with tf.GradientTape() as tape:
   y = self.model.predict(states[t])
   y = tf.reshape(y, constants.NUM_ACTIONS)
   loss = y * mask[t]