compare staging vs FIFO

benchmark_staging.py

import tensorflow as tf
from tensorflow.python.ops import data_flow_ops
import time
from contextlib import contextmanager


@contextmanager
def benchmark(name="unnamed context"):
    elapsed = time.time()
    yield
    elapsed = time.time() - elapsed
    print('[{}] finished in {} ms'.format(name, int(elapsed * 1000)))


tf.reset_default_graph()
with tf.Session() as sess:
    dummy = tf.Variable(tf.random_normal([225, 225, 3], dtype=tf.float32, stddev=1e-1))
    sess.run(tf.global_variables_initializer())
    stager = data_flow_ops.StagingArea([tf.float32])
    enqeue_op = stager.put([dummy])
    dequeue_op = tf.reduce_sum(stager.get())

    for i in range(1000):
        sess.run(enqeue_op)

    with benchmark("staging"):
        s = 0
        for i in range(1000):
            s += sess.run(dequeue_op)


tf.reset_default_graph()
with tf.Session() as sess:
    dummy = tf.Variable(tf.random_normal([225, 225, 3], dtype=tf.float32, stddev=1e-1))
    sess.run(tf.global_variables_initializer())
    stager = tf.FIFOQueue(1000, [tf.float32])
    enqeue_op = stager.enqueue([dummy])
    dequeue_op = tf.reduce_sum(stager.dequeue())

    for i in range(1000):
        sess.run(enqeue_op)

    with benchmark("FIFO"):
        s = 0
        for i in range(1000):
            s += sess.run(dequeue_op)

For those who may find this -- I got the following result. TensorFlow 1.10, Ubuntu, Using a TITAN Xp GPU.

[staging] finished in 396 ms
[FIFO] finished in 579 ms