我用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倍。这是我不明白的,因为理论上第二个因为数据批次的额外序列化/反序列化而应该更慢。
谢谢!
答案 0 :(得分:13)
目前(截至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中提供。在此期间,您可以通过切换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):从评论中回答您的问题:
- 为什么
醇>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
)时间,但解析所涉及的常量比执行操作要小得多。 )
- 为什么要在管道末尾添加预取?
醇>
当您对性能感兴趣时,这几乎总是最好的事情,它应该可以提高管道的整体性能。有关最佳做法的详细信息,请参阅performance guide for tf.data
。