tf.divide() and tf.cast() will interrupt the gradient propagation of my program. Is there any way to solve this problem?

Viewed 380

In python-tensorflow,tf.divide() and tf.cast() will interrupt the gradient propagation of my program. Is there any way to solve this problem? Thank you.

import tensorflow as tf


tf.compat.v1.disable_eager_execution()

w0 = tf.constant([[[0, 1, 2], [1, 2, 0]], [[1, -1, 2], [0, 3, 1]]], dtype=tf.float32)
b0 = tf.constant([[[1, 2, 0], [-1, -2, 1]], [[-2, 1, 0], [0, 1, -1]]], dtype=tf.float32)
w = tf.argmax(w0, axis=-1)
b = tf.argmax(b0, axis=-1)

loss = tf.divide(w, b)
grad1 = tf.gradients(loss, w)
grad2 = tf.gradients(w, w0)
with tf.compat.v1.Session() as sess:
    print(sess.run(loss))
    print(sess.run(grad1))
    print(sess.run(grad2))

Just like this simple example, the program will be in print( sess.run (grad1)) and print( sess.run (grad2)) these two places prompt an error: fetch argument none has invalid type < class' nonetype '>

1 Answers

The problem in your example is not with tf.divide, but with tf.argmax. This function returns an integer value and is not differentiable. You could cast w to float and then compute the gradient with respect to that casted value:

import tensorflow as tf
tf.compat.v1.disable_eager_execution()

w0 = tf.constant([[[0,1,2],[1,2,0]],[[1,-1,2],[0,3,1]]], dtype=tf.float32)
b0 = tf.constant([[[1,2,0],[-1,-2,1]],[[-2,1,0],[0,1,-1]]], dtype=tf.float32)
w = tf.argmax(w0, axis=-1)
b = tf.argmax(b0, axis=-1)

w_f = tf.cast(w, tf.float32)
b_f = tf.cast(b, tf.float32)

loss = tf.divide(w_f, b_f)
grad1 = tf.gradients(loss, w_f)[0]
with tf.compat.v1.Session() as sess:
    print(sess.run(loss))
    # [[2.  0.5]
    #  [2.  1. ]]
    print(sess.run(grad1))
    # [[1.  0.5]
    #  [1.  1. ]]

But you will not be able to compute the gradient of loss with respect to w0, because there is a non-differentiable operation between them.

Related