How to add if condition in a TensorFlow graph?

Question

Let's say I have following code:

x = tf.placeholder("float32", shape=[None, ins_size**2*3], name = "x_input")
condition = tf.placeholder("int32", shape=[1, 1], name = "condition")
W = tf.Variable(tf.zeros([ins_size**2*3,label_option]), name = "weights")
b = tf.Variable(tf.zeros([label_option]), name = "bias")

if condition > 0:
    y = tf.nn.softmax(tf.matmul(x, W) + b)
else:
    y = tf.nn.softmax(tf.matmul(x, W) - b)

Would the if statement work in the calculation (I do not think so)? If not, how can I add an if statement into the TensorFlow calculation graph?

jkschin · Accepted Answer · 2017-07-12 23:18:21Z

109

You're correct that the if statement doesn't work here, because the condition is evaluated at graph construction time, whereas presumably you want the condition to depend on the value fed to the placeholder at runtime. (In fact, it will always take the first branch, because condition > 0 evaluates to a Tensor, which is "truthy" in Python.)

To support conditional control flow, TensorFlow provides the tf.cond() operator, which evaluates one of two branches, depending on a boolean condition. To show you how to use it, I'll rewrite your program so that condition is a scalar tf.int32 value for simplicity:

x = tf.placeholder(tf.float32, shape=[None, ins_size**2*3], name="x_input")
condition = tf.placeholder(tf.int32, shape=[], name="condition")
W = tf.Variable(tf.zeros([ins_size**2 * 3, label_option]), name="weights")
b = tf.Variable(tf.zeros([label_option]), name="bias")

y = tf.cond(condition > 0, lambda: tf.matmul(x, W) + b, lambda: tf.matmul(x, W) - b)

edited Jul 12, 2017 at 23:18

jkschin

5,8646 gold badges38 silver badges62 bronze badges

answered Mar 6, 2016 at 22:00

mrry

126k27 gold badges404 silver badges401 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

Piotr Dabkowski Over a year ago

@mrry Are both branches executed by default ? I have tf.cond(c, lambda x: train_op1, lambda x: train_op2) and both train_ops are executed at each execution of cond independently of the value of c. Am I doing something wrong?

mrry Over a year ago

@PiotrDabkowski This is a sometimes surprising behavior of tf.cond(), which is touched upon in the docs. In short, you need to create the ops that you want to run conditionally inside the respective lambdas. Everything that you create outside the lambdas but refer to in either branch will execute in both cases.

Piotr Dabkowski Over a year ago

@mrry Wow, that's rather unexpected :) Thanks for the answer, defining ops inside functions solved the problem.

AB_IM Over a year ago

Is the condition/( application of the logic) element-wise?

cs95 · Accepted Answer · 2019-06-16 13:03:24Z

11

TensorFlow 2.0

TF 2.0 introduces a feature called AutoGraph which lets you JIT compile python code into Graph executions. This means you can use python control flow statements (yes, this includes if statements). From the docs,

AutoGraph supports common Python statements like while, for, if, break, continue and return, with support for nesting. That means you can use Tensor expressions in the condition of while and if statements, or iterate over a Tensor in a for loop.

You will need to define a function implementing your logic and annotate it with tf.function. Here is a modified example from the documentation:

import tensorflow as tf

@tf.function
def sum_even(items):
  s = 0
  for c in items:
    if tf.equal(c % 2, 0): 
        s += c
  return s

sum_even(tf.constant([10, 12, 15, 20]))
#  <tf.Tensor: id=1146, shape=(), dtype=int32, numpy=42>

edited Jun 16, 2019 at 13:03

answered Jun 16, 2019 at 5:36

cs95

406k106 gold badges744 silver badges797 bronze badges

2 Comments

problemofficer - n.f. Monica Over a year ago

Why are you using tf.equal()? Shouldn't you be able to use == and let AutoGraph compile it automatically?

cs95 Over a year ago

@problemofficer That's a gooooood question. I assumed (like you) the same thing but got bitten. Here is a question I asked which discusses this behavior: stackoverflow.com/questions/56616485/…

Collectives™ on Stack Overflow

How to add if condition in a TensorFlow graph?

2 Answers 2

4 Comments

TensorFlow 2.0

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

4 Comments

TensorFlow 2.0

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related