使用Tensorflow Dataset API读取TFRecords文件时，预处理输入数据会减慢输入管道的速度

Question

I am using Tensorflow Dataset API to read TFRecords files, but the GPU usage is still low (10%). I reckon the cause is that I preprocess the data before they are fed into the sess.run() . Here is my code below.
1. Create a dataset from 3 separate files.

tf.reset_default_graph()

# The content of TFRecords files is that each row is an array. Calculate total rows.
n_total_row = sum(1 for _ in tf.python_io.tf_record_iterator(epd))

def get_epd_dataset(filename):
    dataset = tf.data.TFRecordDataset(filename)
    def _parse_function(example_proto):
        keys_to_features = {'data':tf.VarLenFeature(tf.int64)}
        parsed_features = tf.parse_single_example(example_proto, keys_to_features)
    return tf.sparse_tensor_to_dense(parsed_features['data'])
    # Parse the record into tensors.
    dataset = dataset.map(_parse_function)
    return dataset

# There are 3 essential files comprising input data. It reads 3 seperate
# files "epd", "y_id", "x_feat" into 3 separate dataset respectively, and 
# uses `Dataset.zip()` to combine these 3 separate files into 1 dataset.
epd_ds = get_epd_dataset(epd)
n_lexicon, id_ds = get_id_dataset(y_id)
feat_ds = get_feat_dataset(x_feat)
data_ds = tf.data.Dataset.zip((feat_ds, epd_ds, id_ds))

# Shuffle the dataset
data_ds = data_ds.shuffle(buffer_size=n_total_row, reshuffle_each_iteration=True)
# Repeat the input indefinitly
data_ds = data_ds.repeat(epoch)
# Generate batches
data_ds = data_ds.batch(1)
# Create a one-shot iterator
iterator = data_ds.make_one_shot_iterator()
data_iter = iterator.get_next()

2. Build a Tensorflow graph.

n_input = DIM*(LEFT+1+RIGHT)
n_classes = n_lexicon

mlp = MultiLayerPerceptron.MultiLayerPerceptron(DIM*(LEFT+1+RIGHT), n_lexicon)
# tf Graph input
X = tf.placeholder("float", [None, n_input])
Y = tf.placeholder("float", [None, n_classes])
logits = mlp.multilayer_perceptron(X, dropout_mode)
loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=Y), name='loss_op')
optimizer = tf.train.AdamOptimizer(learning_rate=lr)
train_op = optimizer.minimize(loss_op, name='train_op')

3. Generate data from data_iter and run TF session.

sess = tf.Session()
# Initialization
sess.run(tf.global_variables_initializer())
for e in range(1, epoch+1):
    while True:
        try:
            # Get data from dataset iterator 
            tmp = sess.run([data_iter])[0]
            # a,b,c are a row from 3 serapate files.
            a = tmp[0].flatten()
            b = tmp[1].flatten()
            c = tmp[2].flatten()

            # I believe this step slows down my input pipeline.
            x_train, y_train = _data_generate(mlp, b, d, c)
            _, c = sess.run([train_op, loss_op], feed_dict={X: x_train,
                                                            Y: y_train})
        except tf.errors.OutOfRangeError:
            break
sess.close() 

My code reaches about 10~15% of GPU usage. I think the cause is that _data_generate() consumes too much time on processing numpy array. But I don't know how to improve my pipeline. Here are my questions.

1. According to Tensorflow performance guide and Importing Data , I think using Dataset API and TFRecords files is my best option to solve this low-GPU-usage problem. Or should I use python multithread to feed data into a buffer first and then feed data to sess.run() . I didn't choose the latter solution due to this website mention that

We found that using tf.FIFOQueue and tf.train.queue_runner could not saturate multiple current generation GPUs when using large inputs and processing with higher samples per second,

2. I think that putting _data_generate() in _parse_function() may solve this problem, bucause Tensorflow handles preprocessing data part but not python. But I don't know how to do this since _data_generate() needs 3 rows from 3 separate files. Does anyone know how to do this?

3. Are there other methods could solve my low-GPU-usage problem?

Thank you.

Answer 1

Can you share the code of _data_generate function? I can't see what it does.

As you pointed out performance is likely lost because of RAM <-> GPU memory swap and mixing tensorflow ops with pythonic ones.

Instead of running iterator data_iter yourself by sess.run() , doing numpy operations and then training step, pass data_iter as input to your neural network graph - it should replace the placeholders. (just make a function that constructs the graph using data_iter as parameter).

I think that putting _data_generate() in _parse_function() may solve his problem, bucause Tensorflow handles preprocessing data part but not >python. But I don't know how to do this since _data_generate() needs 3 >rows from 3 separate files. Does anyone know how to do this?

The proper way is to create 3 datasets from files, decode them, zip them, and then pass the iterator to zipped dataset as input to processing graph. You're almost doing that.

Also; Try to enforce multithreading whenever it is possible/needed. Here:

...
return tf.sparse_tensor_to_dense(parsed_features['data'])
    # Parse the record into tensors.
    dataset = dataset.map(_parse_function)
    return dataset

You should use:

dataset.map(_parse_function, num_threads=<MORE THAN ONE>)

Where <MORE THAN ONE> is an integer bigger than one. In your case I would start with 8 threads (see if GPU will be 100%)

Check dis out and tell me if its ok

Answer 2

I'm assuming your example uses a simplified version of your model, otherwise the GPU will almost always terminate its work before the next batch is ready.

Each dataset and transofrmation pipeline has its own specificities and it's difficult to provide a definite answer, but here might be some points worth investigating:

• You are not supposed to compute the values of data_iter explicitely, and you should not use placeholders and feed_dict anymore. The values returned by .make_one_shot_iterator().get_next() are already the input nodes of your model, not the placeholder variables. This will glue together the input pipeline with your model and avoid sending data back and forth between tensorflow CPU, python and tensorflow GPU memory regions. In your session.run() call, you need not specify any input value, the model will automatically feed itself on the dataset as needed.
• I don't think TFRecord files are supposed to be fast on random access due to the unknown element size. In your case, you also seem to use random access three times: for epd, feats and the metadata.
• Did you instruct tensorflow to use prefetching somewhere? Otherwise tensorflow may wait for the current batch to be processed before loading and processing the next one.
• You might want to check if the multiple workers of the map function are not waiting for some resource (disk IO, CPU time, locked object?)
• tf.data API is tricky to debug and not very practical (seriously, how could it not know about anything about tensors?), you may want to have a look at the tensorpack library or the old tensorflow Queue data loading API.