我创建了一个 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)。
在第二种方法中,我手动将数据集中的批处理放入 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 倍。这是我不明白的,因为理论上第二个应该会因为数据批处理的额外序列化/反序列化而变慢。
谢谢!
最佳答案
当前(从 TensorFlow 1.9 开始)使用 tf.data
时存在性能问题映射和批处理具有大量特征的张量,每个张量中的数据量很少。该问题有两个原因:
dataset.map(parse_tfrecord, ...)
转换将执行 O( batch_size
* num_columns
) 小操作来创建批处理。相比之下,喂食 tf.placeholder()
至tf.parse_example()
将执行 O(1) 操作来创建相同的批处理。 tf.SparseTensor
使用 dataset.batch()
的对象比直接创建相同的 tf.SparseTensor
慢得多作为 tf.parse_example()
的输出. 对这两个问题的改进正在进行中,并且应该在 TensorFlow 的 future 版本中可用。同时,您可以提高
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) : 要从评论中回答您的问题:
- Why is
dataset.map(parse_tfrecord, ...)
O(batch_size
*num_columns
), not O(batch_size
)? If parsing requires enumeration of the columns, why doesn't parse_example take O(num_columns
)?
当您将 TensorFlow 代码包装在
Dataset.map()
中时(或其他函数转换)每个输出的恒定数量的额外操作被添加到函数的“返回”值中,并且(在 tf.SparseTensor
值的情况下)将它们“转换”为标准格式。当您直接传递 tf.parse_example()
的输出时对于模型的输入,不会添加这些操作。虽然它们是非常小的操作,但执行如此多的操作可能会成为瓶颈。 (从技术上讲,解析确实需要 O( batch_size
* num_columns
) 时间 ,但是解析中涉及的常量比执行操作要小得多。)
- Why do you add a prefetch at the end of the pipeline?
当您对性能感兴趣时,这几乎总是最好的做法,它应该会提高管道的整体性能。有关最佳实践的更多信息,请参阅 performance guide for
tf.data
.
关于tensorflow - 使用 feed_dict 比使用数据集 API 快 5 倍以上?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/50781373/