Tensorflow: Custom Operation Used In Two Networks Simultaneously Produces Nan

I have written following custom operation with gradient to binarize a real vector. (this code is inspired from https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342) def

Solution 1:

Try this maybe.

@function.Defun()defBinarizerGrad(unused_x, dy):
  # Backprop dy directly.return dy

@function.Defun(grad_func=BinarizerGrad)defBinarizer(x):
  # your whatever forward function here. return tf.floor(x + tf.random_uniform(tf.shape(x)))

g = tf.Graph()
with g.as_default():
  x = tf.placeholder(tf.float32)
  y = Binarizer(x)
  dy = tf.placeholder(tf.float32)
  dx = tf.gradients(y, x, grad_ys=dy)

with tf.Session(graph=g) as sess:
  x_val = np.array([[1, 2, 3], [0.5, 0.3, 0.2]])
  dy_val = np.array([[1, 0, 1], [0., 0.1, 0.9]])
  for v in sess.run([x, y, dx], feed_dict={x : x_val, dy: dy_val}):
    print v

Solution 2:

I don't think building a graph, starting a session, and running is well-supported inside a py_func. In this case you can remove all those things and just use straight-up tensorflow code and everything should work.