4

I'm changing my TensorFlow code from the old queue interface to the new Dataset API. With the old interface I could monitor the actual filled queue size by accessing a raw counter in the graph, e.g. as follows:

queue = tf.train.shuffle_batch(...,  name="training_batch_queue")
queue_size_op = "training_batch_queue/random_shuffle_queue_Size:0"
queue_size = session.run(queue_size_op)

However, with the new Dataset API I can't seem to find any variables in the graph related to the queues / datasets, so my old code doesn't work anymore. Is there any way to obtain the number of items in the queue using the new Dataset API (e.g. in the tf.Dataset.prefetch or tf.Dataset.shuffle queue)?

It is important for me to monitor the number of items in the queue, as that tells me a lot about the behaviour of the pre-processing in the queues, including whether the pre-processing or the remainder (e.g. a neural network) is the speed bottleneck.

CNugteren
  • 860
  • 9
  • 14
  • please, have a look at [46444018](https://stackoverflow.com/questions/46444018/meaning-of-buffer-size-in-dataset-map-dataset-prefetch-and-dataset-shuffle) to get a better idea of the underlying behavior of the different kinds of shuffle arguments – Max F. Feb 27 '18 at 14:14
  • @maxF. Yes, I understand. The example in my post is perhaps not the best, since it is interesting to monitor `tf.train.shuffle_batch` in the old setting, but it makes no sense to monitor `tf.Dataset.shuffle` in the new setting. What does make sense to monitor is the size of `tf.Dataset.prefetch`, to get an idea of whether the pre-processing or the actual network is the bottleneck. – CNugteren Mar 01 '18 at 13:02

1 Answers1

0

As a work around it is possible to keep a counter to indicate how many items are in the queue. Here's how to define the counter:

 queue_size = tf.get_variable("queue_size", initializer=0,
                              trainable=False, use_resource=True)

Then, when pre-processing data (e.g. in the dataset.map function), we can increment that counter:

 def pre_processing():
    data_size = ... # compute this (could be just '1')
    queue_size_op = tf.assign_add(queue_size, data_size)  # adding items
    with tf.control_dependencies([queue_size_op]):
        # do the actual pre-processing here

We can then decrement the counter every-time we run our model with a batch of data:

 def model():
    queue_size_op = tf.assign_add(queue_size, -batch_size)  # removing items
    with tf.control_dependencies([queue_size_op]):
        # define the actual model here

Now, all we need to do is run the queue_size tensor in our training loop to find out what the current queue size is, i.e. the number of items in the queue at this moment:

 current_queue_size = session.run(queue_size)

It's a bit less elegant compared to the old way (before the Dataset API), but it does the trick.

CNugteren
  • 860
  • 9
  • 14