How to properly apply gradients in TensorFlow v2.0

user_7091135 0
Viewed 737

I'm trying to reproduce some neural transfer tutorial which previously based on TensorFlow v1.x but using TensorFlow v2.0. For some reason, I can't apply gradients with Adam optimizer due to, I believe, some data format. The source is quite large so I tried to give just most relevant lines, but if someone asks about full code I will add it here.

Here is the optimizer

    opt = tf.optimizers.Adam(learning_rate = 5, beta_1 = 0.99, epsilon = 1e-1)

And when I'm trying to apply gradients to initial variables using

    opt.apply_gradients(zip(grads, init_image))

I'm getting error

!opt.apply_gradients(zip(grads, init_image))
*** TypeError: zip argument #2 must support iteration

These are the zip arguments:

init_image
<tf.Variable 'Variable:0' shape=(1, 345, 512, 3) dtype=float32, numpy=
array([[[[...]]]], dtype=float32)>

grads
<tf.Tensor: shape=(1, 345, 512, 3), dtype=float32, numpy=
array([[[[...]]]], dtype=float32)>

As you can see, the shape is the same as it should be, values inside numpy are some real values, not nans or something like that.

The gradients are calculated with:

    with tf.GradientTape() as tape: 
         ...computing all_loss...
    total_loss = all_loss[0]
    grads = tape.gradient(total_loss, init_image)

Any suggestions please.

tensorflow2.0 tensorflow machine-learning numpy python
1 Answers

This should be rather simple to fix -- the issue is that you need to pack the gradients/variables in a list for zip to work. So simply use

opt.apply_gradients(zip([grads], [init_image]))

To add some explanation -- your code might run without crashing in eager mode, but it could produce incorrect results (would have to test this). It definitely crashes when using tf.function (whether you are using this is not clear from the code you provide) because tensors are not iterable in graph mode. By wrapping them in lists, zip can now iterate over those (even though there's only one element).

user_93931020
Related