使用feed_dict比使用数据集API快5倍？

Question

我创建了一个TFRecord格式的数据集进行测试.每个条目包含200个列,名为C1- C199,每个都是字符串列表,以及一个label用于表示标签的列.创建数据的代码可以在这里找到:https://github.com/codescv/tf-dist/blob/8bb3c44f55939fc66b3727a730c57887113e899c/src/gen_data.py#L25

然后我使用线性模型来训练数据.第一种方法如下:

dataset = tf.data.TFRecordDataset(data_file)
dataset = dataset.prefetch(buffer_size=batch_size*10)
dataset = dataset.map(parse_tfrecord, num_parallel_calls=5)
dataset = dataset.repeat(num_epochs)
dataset = dataset.batch(batch_size)

features, labels = dataset.make_one_shot_iterator().get_next()    
logits = tf.feature_column.linear_model(features=features, feature_columns=columns, cols_to_vars=cols_to_vars)
train_op = ...

with tf.Session() as sess:
    sess.run(train_op)

完整的代码可以在这里找到:https://github.com/codescv/tf-dist/blob/master/src/lr_single.py

当我运行上面的代码时,我得到0.85步/秒(批量大小为1024).

在第二种方法中,我手动从Dataset获取批处理到python,然后将它们提供给占位符,如下所示:

example = tf.placeholder(dtype=tf.string, shape=[None])
features = tf.parse_example(example, features=tf.feature_column.make_parse_example_spec(columns+[tf.feature_column.numeric_column('label', dtype=tf.float32, default_value=0)]))
labels = features.pop('label')
train_op = ...

dataset = tf.data.TFRecordDataset(data_file).repeat().batch(batch_size)
next_batch = dataset.make_one_shot_iterator().get_next()

with tf.Session() as sess:
    data_batch = sess.run(next_batch)
    sess.run(train_op, feed_dict={example: data_batch})

完整的代码可以在这里找到:https://github.com/codescv/tf-dist/blob/master/src/lr_single_feed.py

当我运行上面的代码时,我得到5步/秒.这比第一种方法快5倍.这是我不明白的,因为理论上第二个应该是因为数据批量的额外序列化/反序列化而变慢.

谢谢!

Answer 1

目前(截至TensorFlow 1.9)在使用tf.data映射和批量张量的张量时存在性能问题,这些张量具有大量的特征,每个特征中都有少量数据.这个问题有两个原因:

1. 该dataset.map(parse_tfrecord, ...)转型将执行O(batch_size*num_columns)小操作创建一个批处理.相反,馈送tf.placeholder()到tf.parse_example()将执行O(1)操作以创建相同的批次.

2. 使用批处理tf.SparseTensor对象dataset.batch()比直接创建与tf.SparseTensor输出相同的速度慢得多tf.parse_example().

对这两个问题的改进正在进行中,并且应该在未来版本的TensorFlow中提供.在此期间,可以提高性能tf.data通过切换的顺序为基础的管道dataset.map()和dataset.batch()和重写dataset.map()到字符串矢量工作,如基于喂养版本:

dataset = tf.data.TFRecordDataset(data_file)
dataset = dataset.prefetch(buffer_size=batch_size*10)
dataset = dataset.repeat(num_epochs)

# Batch first to create a vector of strings as input to the map(). 
dataset = dataset.batch(batch_size)

def parse_tfrecord_batch(record_batch):
  features = tf.parse_example(
      record_batch,
      features=tf.feature_column.make_parse_example_spec(
          columns + [
              tf.feature_column.numeric_column(
                  'label', dtype=tf.float32, default_value=0)]))
  labels = features.pop('label')
  return features, labels

# NOTE: Parallelism might not be as useful, because the individual map function now does
# more work per invocation, but you might want to experiment with this.
dataset = dataset.map(parse_tfrecord_batch)

# Add a prefetch at the end to pipeline execution.
dataset = dataset.prefetch(1)

features, labels = dataset.make_one_shot_iterator().get_next()    
# ...

编辑(2018/6/18):从评论中回答您的问题:

1. 为什么是dataset.map(parse_tfrecord, ...)O(batch_size*num_columns),而不是O(batch_size)？如果解析需要枚举列,为什么parse_example不需要O(num_columns)？

当您将TensorFlow代码包装在Dataset.map()(或其他功能转换)中时,每个输出的常量数量的额外操作被添加到函数的"返回"值,并且(在tf.SparseTensor值的情况下)将它们"转换"为标准格式.直接将输出传递tf.parse_example()给模型的输入时,不会添加这些操作.虽然它们是非常小的操作,但执行这么多操作可能会成为瓶颈.(从技术上讲,解析确实需要O(batch_size*num_columns)时间,但解析中涉及的常量比执行操作要小得多.)

2. 为什么要在管道末尾添加预取？

当您对性能感兴趣时,这几乎总是最好的事情,它应该可以提高管道的整体性能.有关最佳实践的详细信息,请参阅性能指南tf.data.