Error while calculating gradient penalty for wgan-gp

Question

i was calculating loss function in wgan-gp and I wanted to know what was wrong with my code or if some other method needs to be implemented

with tf.GradientTape() as critic_tape:
     generated_images = generator(tf.random_normal([16, 100]), training=True)
     a = tf.convert_to_tensor(images[:16])
     real_output = critic(a, training=True)
    generated_output = critic(generated_images, training=True)               
    with tf.GradientTape() as gtape:
        epsilon = tf.random_uniform([], 0, 1)
        xhat = epsilon*a + (1-epsilon)*generated_images
        dhat = critic(xhat, training=True)
        gtape.watch(xhat)
   dhat2 = gtape.gradient(dhat, xhat)
   slopes = tf.sqrt(tf.reduce_sum(tf.square(dhat2), reduction_indices=[1]))
   gradient_penalty = 10*tf.reduce_mean((slopes-1.0)**2)
critic_loss = get_critic_loss(real_output, generated_output)
critic_loss+= gradient_penalty                
gradients_of_critic = critic_tape.gradient(critic_loss, critic.variables)

this is the error stack, i am using tensorflow eager execution, so any help would be really appreciated

---------------------------------------------------------------------------
LookupError                               Traceback (most recent call last)
<ipython-input-512-cbc8ebf905ac> in <module>()
     16     critic_loss = get_critic_loss(real_output, generated_output)
     17     critic_loss+= gradient_penalty
---> 18 gradients_of_critic = critic_tape.gradient(critic_loss, critic.variables)
     19 print(gradients_of_critic)

c:\users\vibhu\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\backprop.py in gradient(self, target, sources, output_gradients)
    856     flat_grad = imperative_grad.imperative_grad(
    857         _default_vspace, self._tape, nest.flatten(target), flat_sources,
--> 858         output_gradients=output_gradients)
    859 
    860     if not self._persistent:

c:\users\vibhu\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\imperative_grad.py in imperative_grad(vspace, tape, target, sources, output_gradients)
     61   """
     62   return pywrap_tensorflow.TFE_Py_TapeGradient(
---> 63       tape._tape, vspace, target, sources, output_gradients)  # pylint: disable=protected-access

c:\users\vibhu\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\backprop.py in _gradient_function(op_name, attr_tuple, num_inputs, inputs, outputs, out_grads)
    110   """
    111   mock_op = _MockOp(attr_tuple, inputs, outputs, op_name)
--> 112   grad_fn = ops._gradient_registry.lookup(op_name)  # pylint: disable=protected-access
    113   if grad_fn is None:
    114     return [None] * num_inputs

c:\users\vibhu\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\framework\registry.py in lookup(self, name)
     91     else:
     92       raise LookupError(
---> 93           "%s registry has no entry for: %s" % (self._name, name))

LookupError: gradient registry has no entry for: StatefulPartitionedCall

Answer 1

i faced the same issue, likely building on the same code as you. Disabling tf.contrib.eager.defun for my critic solved my issue. This is of course not a good solution, because you wont be getting the defun speedup but at least you can execute your code.