CODEWITHSUNDEEP
pythonpython-3.xtensorflow
Hoeze
Hoeze

Reputation: 716

tf.data.Iterator.get_next(): How to advance in tf.while_loop?

Currently I try to implement all training in a Tensorflow while loop, but I've got problems with the Tensorflow dataset API's Iterator.

Usually, when calling sess.run(), Iterator.get_next() advances to the next element. However, I need to advance to the next element INSIDE one run. How do I do this?

The following small example shows my problem:

import tensorflow as tf
import numpy as np


def for_loop(condition, modifier, body_op, idx=0):
    idx = tf.convert_to_tensor(idx)

    def body(i):
        with tf.control_dependencies([body_op(i)]):
            return [modifier(i)]

    # do the loop:
    loop = tf.while_loop(condition, body, [idx])
    return loop


x = np.arange(10)

data = tf.data.Dataset.from_tensor_slices(x)
data = data.repeat()

iterator = data.make_initializable_iterator()
smpl = iterator.get_next()

loop = for_loop(
    condition=lambda i: tf.less(i, 5),
    modifier=lambda i: tf.add(i, 1),
    body_op=lambda i: tf.Print(smpl, [smpl], message="This is sample: ")
)

sess = tf.InteractiveSession()
sess.run(iterator.initializer)
sess.run(loop)

Output:

This is sample: [0]
This is sample: [0]
This is sample: [0]
This is sample: [0]
This is sample: [0]

I always get exactly the same element.

Upvotes: 3

Views: 1622

Answers (1)

benjaminplanche
benjaminplanche

Reputation: 15119

You need to call iterator.get_next() every time you want to "iterate inside one run".

For instance in your toy example, just replace your body_op with:

 body_op=lambda i: tf.Print(i, [iterator.get_next()], message="This is sample: ")
# This is sample: [0]
# This is sample: [1]
# This is sample: [2]
# This is sample: [3]
# This is sample: [4]

Upvotes: 2

Related Questions