Tensorflow while_loop for training

Question

In my problem I need run GD with 1 example from data on each training step. It's known problem that session.run() has overhead and therefore it is too long to train model. In attempt to avoid overhead I tried to use while_loop and train model on all data with one run() call. But it approach don't work and train_op don't execute even ones. Below simple example of what I'm doing:

data = [k*1. for k in range(10)]
tf.reset_default_graph()

i = tf.Variable(0, name='loop_i')
q_x = tf.FIFOQueue(100000, tf.float32)
q_y = tf.FIFOQueue(100000, tf.float32)

x = q_x.dequeue()
y = q_y.dequeue()
w = tf.Variable(0.)
b = tf.Variable(0.)
loss = (tf.add(tf.mul(x, w), b) - y)**2

gs = tf.Variable(0)

train_op = tf.train.GradientDescentOptimizer(0.05).minimize(loss, global_step=gs)

s = tf.Session()
s.run(tf.initialize_all_variables())

def cond(i):
    return i < 10

def body(i):
    return tf.tuple([tf.add(i, 1)], control_inputs=[train_op])


loop = tf.while_loop(cond, body, [i])

for _ in range(1):
    s.run(q_x.enqueue_many((data, )))
    s.run(q_y.enqueue_many((data, )))

s.run(loop)
s.close()

What I'm doing wrong? Or there is another solution of this problem with too expensive overhead?

Thanks!

Answer 1

The reason the model does not appear to train is because the input reading, gradient calculation, and the minimize() call are all defined outside (and hence, in dataflow terms, before) the body of the tf.while_loop(). This means that all of these parts of the model run only once, before the loop executes, and the loop itself has no effect.

A slight refactoring—to move the dequeue() operations, gradient calculation, and minimize() call inside the loop—fixes the problem and allows your program to train:

optimizer = tf.train.GradientDescentOptimizer(0.05)

def cond(i):
    return i < 10

def body(i):
    # Dequeue a new example each iteration.
    x = q_x.dequeue()
    y = q_y.dequeue()

    # Compute the loss and gradient update based on the current example.
    loss = (tf.add(tf.mul(x, w), b) - y)**2
    train_op = optimizer.minimize(loss, global_step=gs)

    # Ensure that the update is applied before continuing.
    return tf.tuple([tf.add(i, 1)], control_inputs=[train_op])

loop = tf.while_loop(cond, body, [i])

---

UPDATE: Here's a complete program the executes the while loop, based on the code in your question:

import tensorflow as tf

# Define a single queue with two components to store the input data.
q_data = tf.FIFOQueue(100000, [tf.float32, tf.float32])

# We will use these placeholders to enqueue input data.
placeholder_x = tf.placeholder(tf.float32, shape=[None])
placeholder_y = tf.placeholder(tf.float32, shape=[None])
enqueue_data_op = q_data.enqueue_many([placeholder_x, placeholder_y])

gs = tf.Variable(0)
w = tf.Variable(0.)
b = tf.Variable(0.)
optimizer = tf.train.GradientDescentOptimizer(0.05)

# Construct the while loop.
def cond(i):
    return i < 10

def body(i):
    # Dequeue a single new example each iteration.
    x, y = q_data.dequeue()
    # Compute the loss and gradient update based on the current example.
    loss = (tf.add(tf.multiply(x, w), b) - y) ** 2
    train_op = optimizer.minimize(loss, global_step=gs)
    # Ensure that the update is applied before continuing.
    with tf.control_dependencies([train_op]):
        return i + 1

loop = tf.while_loop(cond, body, [tf.constant(0)])

data = [k * 1. for k in range(10)]

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for _ in range(1):
        # NOTE: Constructing the enqueue op ahead of time avoids adding
        # (potentially many) copies of `data` to the graph.
        sess.run(enqueue_data_op,
                 feed_dict={placeholder_x: data, placeholder_y: data})
    print (sess.run([gs, w, b]))  # Prints before-loop values.
    sess.run(loop)
    print (sess.run([gs, w, b]))  # Prints after-loop values.