Example of pipelining with StagingArea with tf.tuple()

staging_area_pipelining_with_tf_tuple.py

# Example of pipelining with StagingArea with tf.tuple().
#
# In particular it shows a trick how to group together get() and put()
# operations using tf.tuple() such that we have a single operation that
# works like# semicolon operator: first it performs put(), then get()
# and returns the tensor output of get().
#
# We could possibly use that compound operation as Keras model input
# so that we don't need to modify K.function() to pass additional
# fetches to Session.run() explicitly. However it's not sure if this
# mechanism will really allow HtoD memcpy in parallel to computation.
#
# More info: https://www.tensorflow.org/api_docs/python/tf/tuple

import numpy as np
import tensorflow as tf
from tensorflow.python.ops.data_flow_ops import StagingArea

length = 10
dataset_range = tf.contrib.data.Dataset.range(length)
iter = dataset_range.make_one_shot_iterator()
next_item = iter.get_next()

area = StagingArea(dtypes=[tf.int64])
area_put = area.put([next_item])
area_get = area.get()
area_size = area.size()
area_get_put = tf.tuple([area_get], control_inputs=[area_put])[0]

with tf.Session() as sess:
    # first item, just put()
    print('put:', sess.run(area_put))
    # get() & put()
    for i in range(length - 1):
        # this works as "semicolon"
        print('put(); get() =', sess.run(area_get_put))
        print('size:', sess.run(area_size))
    # last item, just get()
    print('get() =', sess.run(area_get))
    print('size:', sess.run(area_size))
    
# Output:
# put: None
# put(); get() = 0
# size: 1
# put(); get() = 1
# size: 1
# put(); get() = 2
# size: 1
# put(); get() = 3
# size: 1
# put(); get() = 4
# size: 1
# put(); get() = 5
# size: 1
# put(); get() = 6
# size: 1
# put(); get() = 7
# size: 1
# put(); get() = 8
# size: 1
# get() = 9
# size: 0

Hi, I have a trouble with replacing iter with tf.Variable(). Can you make an example of this? thx