Understanding the while loop in Tensorflow

Question

I am using the Python API for Tensorflow. I am trying to implement the Rosenbrock function given below without the use of a Python loop:

My current implementation is as follows:

def rosenbrock(data_tensor):
    columns = tf.unstack(data_tensor)

    summation = 0
    for i in range(1, len(columns) - 1):
        first_term = tf.square(tf.subtract(columns[i + 1], tf.square(columns[i])))
        second_term = tf.square(tf.subtract(columns[i], 1.0))
        summation += tf.add(tf.multiply(100.0, first_term), second_term)

    return summation

I have tried implementing the summation in a tf.while_loop(); however, I found the API somewhat unintuitive when it comes to using an index integer that is meant to remain separate from the data. The example given in the documentation uses the data as the index (or vice-versa):

i = tf.constant(0)
c = lambda i: tf.less(i, 10)
b = lambda i: tf.add(i, 1)
r = tf.while_loop(c, b, [i])

Is it appropriate to just use the for loop? What's the benefit of employing while_loop? Or is it necessary? — Lerner Zhang, Aug 11 '17 at 08:43
In his code above, the for loop will execute python code. If we call the body of his for loop "f", then you can think of the python code as executing f,f,f,f,f,... f. So it will call that "body" function N times, and the graph of the function will thus have that function N times. If you use a tf.while_loop, then you will only see that function once in the graph. — bremen_matt, Jan 23 '18 at 11:29
The advantage of the tf.while_loop is: 1) you can run iterations in parallel and 2) you can have runtime constants in your condition statements. For instance, if you wanted to run an optimizer until a certain tolerance is met, then you have to use the tf.while_loop variant because python cannot initially evaluate the condition — bremen_matt, Jan 23 '18 at 11:33
@bremen_matt, If iterations have sequential dependencies then how parallel computation works? — amin__, Feb 12 '18 at 04:30
Check this out. I found it very helpful https://github.com/tensorflow/tensorflow/issues/1984 — bremen_matt, Feb 12 '18 at 05:13

score 16 · Accepted Answer · answered May 04 '17 at 21:21

This can be achieved using the tf.while_loop() and standard tuples as per the second example in the documentation.

def rosenbrock(data_tensor):
    columns = tf.unstack(data_tensor)

    # Track both the loop index and summation in a tuple in the form (index, summation)
    index_summation = (tf.constant(1), tf.constant(0.0))

    # The loop condition, note the loop condition is 'i < n-1'
    def condition(index, summation):
        return tf.less(index, tf.subtract(tf.shape(columns)[0], 1))

    # The loop body, this will return a result tuple in the same form (index, summation)
    def body(index, summation):
        x_i = tf.gather(columns, index)
        x_ip1 = tf.gather(columns, tf.add(index, 1))

        first_term = tf.square(tf.subtract(x_ip1, tf.square(x_i)))
        second_term = tf.square(tf.subtract(x_i, 1.0))
        summand = tf.add(tf.multiply(100.0, first_term), second_term)

        return tf.add(index, 1), tf.add(summation, summand)

    # We do not care about the index value here, return only the summation
    return tf.while_loop(condition, body, index_summation)[1]

It is important to note that the index increment should occur in the body of the loop similar to a standard while loop. In the solution given, it is the first item in the tuple returned by the body() function.

Additionally, the loop condition function must allot a parameter for the summation although it is not used in this particular example.

Understanding the while loop in Tensorflow

1 Answers1

Linked