def _set_up_staging(self, transition):
"""Sets up staging ops for prefetching the next transition.
This allows us to hide the py_func latency. To do so we use a staging area
to pre-fetch the next batch of transitions.
Args:
transition: tuple of tf.Tensors with shape
memory.get_transition_elements().
Returns:
prefetched_transition: tuple of tf.Tensors with shape
memory.get_transition_elements() that have been previously prefetched.
"""
transition_type = self.memory.get_transition_elements()
# Create the staging area in CPU.
prefetch_area = StagingArea(
[shape_with_type.type for shape_with_type in transition_type])
# Store prefetch op for tests, but keep it private -- users should not be
# calling _prefetch_batch.
self._prefetch_batch = prefetch_area.put(transition)
initial_prefetch = tf.cond(tf.equal(prefetch_area.size(), 0),
lambda: prefetch_area.put(transition),
tf.no_op)
# Every time a transition is sampled self.prefetch_batch will be
# called. If the staging area is empty, two put ops will be called.
with tf.control_dependencies([self._prefetch_batch, initial_prefetch]):
prefetched_transition = prefetch_area.get()
return prefetched_transition
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
length = 10
dataset_range = tf.data.Dataset.range(length)
iter = dataset_range.make_one_shot_iterator()
next_item = iter.get_next()
with tf.device('cpu:0'):
a = tf.Variable([1], dtype=tf.int64)
area = StagingArea(dtypes=[tf.int64])
area_put = area.put([next_item])
area_get = area.get()[0]
area_size = area.size()
area_get_put = tf.tuple([area_get], control_inputs=[area_put])[0]
b = a + area_get
c = b + area_get
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
# first item, just put()
print('put:', sess.run(area_put))
# get() & put()
for i in range(length - 1):
# this works as "semicolon"
print(sess.run(c))
print('put(); get() =', sess.run(area_put))
print('size:', sess.run(area_size))