what is the equivalent sytax for `Tensor.grad` in Tensorflow 2.0

user_5865579 0
Viewed 89

In the Pytorch, we can access the gradient of a variable x by

z.grad

What is the same sytax in Tensorflow 2. My goad is to cut the gradient. Here is Pytorch code

if z.grad > 1000:
    z.grad = 10

Can tensorflow 2 apply the same functions?

Thanks

tensorflow2.0 pytorch gradient python
1 Answers

So in TF2, assume we define following variables and optimizer:

import tensorflow as tf
from tensorflow import keras 

opt = tf.keras.optimizers.Adam(learning_rate=0.1)

x = tf.Variable([3.0, 4.0]) 
y = tf.Variable([1.0, 1.0, 1.0, 1.0])
var_list = [x, y]

Then we can get gradients by using tf.GradientTape():

with tf.GradientTape() as tape:
    loss = tf.reduce_sum(x ** 2) + tf.reduce_sum(y) 

grads = tape.gradient(loss, var_list)

Finally we could process the gradients by custom function:

def clip_grad(grad):
    if grad > 1000:
        grad = 10
    return grad

processed_grads = [tf.map_fn(clip_grad, g) for g in grads]

opt.apply_gradients(zip(processed_grads, var_list))

Note you may find the keras optimizers have get_gradients method, but it won't work with eager execution enabled which is default in TF2, if you want use that, then you may have to write code in TF1 fashion

user_135095400
Related