Issue
I'm trying to execute some condition-dependent functions where each function needs to contract tensors differently depending on their shapes, for instance. However, I realised that tf.cond
and tf.case
is executing all functions regardless of the condition. Prepared the following code as an example;
def a():
print("a")
return tf.constant(2)
def b():
print("b")
return tf.constant(3)
def c():
print("c")
return tf.constant(4)
def d():
print("default")
return tf.constant(1)
x = tf.constant(1)
@tf.function
def f():
return tf.case([
(tf.equal(x,1), a),
(tf.equal(x,2), b),
(tf.equal(x,2), c)
], default=d, exclusive=True)
@tf.function
def f1():
def cond3():
return tf.cond(tf.equal(x,2), c, d)
def cond2():
return tf.cond(tf.equal(x,2), b, cond3)
return tf.cond(tf.equal(x,1), a, cond2)
print(f())
print(f1())
# Output:
# a
# b
# c
# default
# tf.Tensor(2, shape=(), dtype=int32)
# a
# b
# c
# default
# tf.Tensor(2, shape=(), dtype=int32)
as you can see for both of the cases, the result is as expected but each function is executed while reaching the conclusion. Hence in my particular case, since I'm doing different calculations depending on the tensor's shape, I get a multitude of errors. I've seen many such bug reports but haven't found a solution. Is there another way to do conditional execution that I'm not aware of where different functions can be executed depending on the condition? Note that I tried simply using if tf.equal(x,2): ...
but in that case, I'm getting an error saying that tensor output can not be used as python boolean. Note that this example is much-simplified version of my problem, my conditions are based on tensor shapes such as tf.equal(tf.size(tensor), N)
so I really need a way to execute different things for different cases.
After @LaplaceRicky 's answer I realised that the code that I provided was not representative enough so I'm providing a better example showing what I need to do;
x = tf.ones((3,2,1))
y = tf.ones((1,2,3))
z = tf.ones((4,3,5))
k = tf.ones((3,5,5))
def a(t):
def exe():
return tf.einsum("ijk,lmi", t, y)
return exe
def b(t):
def exe():
return tf.einsum("ijk,ljm", t, z)
return exe
def d(t):
def exe():
return tf.einsum("ijk,klm", t, z)
return exe
c = tf.constant(1)
@tf.function
def f(t):
y = tf.case([
(tf.equal(tf.shape(t)[0], 3), a(t)),
(tf.equal(tf.shape(t)[1], 3), b(t)),
], default=d, exclusive=True)
return y
print(f(x))
This function will execute properly without tf.function
decorator leading to
tf.Tensor(
[[[[3. 3.]]]
[[[3. 3.]]]], shape=(2, 1, 1, 2), dtype=float32
However, when the decorator is included I got a ValueError
which shows that all the cases are executed.
System information
- TensorFlow version: 2.4.1
- Python version: 3.8.2
Solution
Short answer: use tf.print
instead of print
to check whether a particular branch is really being executed in tensorflow graph mode.
Explanations: print
does not work and won't print in graph mode but it will print during tracing. The printed messages actually implies all of the branches were added to the tensorflow graph but it does not imply all branches will be executed all the time in graph mode. tf.print
should be used instead for the debugging.
For more information: https://www.tensorflow.org/guide/function#conditionals
Demonstration:
def a():
tf.print('a')
return tf.constant(10)
def b():
tf.print('b')
return tf.constant(11)
def c():
tf.print('c')
return tf.constant(12)
@tf.function
def cond_fn(x):
return tf.switch_case(x, {0:a,1:b}, default=c)
print(cond_fn(tf.constant(0)))
print(cond_fn(tf.constant(1)))
print(cond_fn(tf.constant(2)))
Expected outputs:
a
tf.Tensor(10, shape=(), dtype=int32)
b
tf.Tensor(11, shape=(), dtype=int32)
c
tf.Tensor(12, shape=(), dtype=int32)
The ValueError
error message is because tensorflow graph does not support this kind of feature very well, at least not with tf.einsum
. One way of the workarounds is to have a graph that supports variable-shaped inputs by using tf.function(f).get_concrete_function(tf.TensorSpec(shape=[None,None,None]))
.
Besides, tf.einsum
is problematic in the process and have to be replaced by tf.transpose
and tf.tensordot
.
Example Codes:
x = tf.random.normal((3,2,1))
y = tf.random.normal((1,2,3))
z = tf.random.normal((4,3,5))
k = tf.random.normal((3,5,5))
#for checking the values
def f2(t):
p = tf.case([
(tf.equal(tf.shape(t)[0], 3), lambda:tf.einsum("ijk,lmi", t, y)),
(tf.equal(tf.shape(t)[1], 3), lambda:tf.einsum("ijk,ljm", t, z)),
], default=lambda:tf.einsum("ijk,klm", t, k), exclusive=True)
return p
#work around
def f(t):
if tf.shape(t)[0] == 3:
tf.print('branch a executed')
return tf.tensordot(tf.transpose(t,[1,2,0]), tf.transpose(y,[2,0,1]),1)
elif tf.shape(t)[1] == 3:
tf.print('branch b executed')
return tf.tensordot(tf.transpose(t,[0,2,1]), tf.transpose(z,[1,0,2]),1)
else:
tf.print('branch c executed')
return tf.tensordot(t, k,1)
graph_f=tf.function(f).get_concrete_function(tf.TensorSpec(shape=[None,None,None]))
print(np.allclose(graph_f(x),f2(x)))
print(np.allclose(graph_f(y),f2(y)))
print(np.allclose(graph_f(z),f2(z)))
Expected outputs:
branch a executed
True
branch c executed
True
branch b executed
True
Answered By - Laplace Ricky
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.