使用Estimators API结合tf.data.Dataset时如何加快批量准备

Question

我想加快我的训练程序,使用带有input_fn编写的Estimator API tf.data.Dataset.

我的实现需要2秒钟来准备一批数据,然后在GPU上运行训练1秒,然后重新开始准备批处理.这实在是效率低下.

我正在寻找一种方法来异步准备批次并将它们上传到GPU以加速培训.或者替代地用于在调用之间缓存数据集的方法input_fn(由于dataset.cache()必须在每个input_fn调用上重新创建数据集,因此似乎不是一个好的选择).

这是我的代码的简化版本:

def input_fn(filenames, labels, epochs):
  dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
  dataset = dataset.map(_read_wav, num_parallel_calls=num_map_threads)
  if shuffle:
     dataset = dataset.shuffle(buffer_size=len(labels))
  dataset = dataset.map(_post_process,  num_parallel_calls=num_map_threads)
  dataset = dataset.map(lambda wav, label: ({'wav': wav}, label))
  dataset = dataset.batch(128)
  dataset = dataset.repeat(epochs) # to iterate over the training set forever
  iterator = dataset.dataset.make_one_shot_iterator()
  features, labels = iterator.get_next()
  return features, labels

train_input_fn = lambda : input_fn(train_files, train_labels, None)
eval_input_fn = lambda : input_fn(eval_files, eval_labels, 1)

train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=45000)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn) 
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

我注意到Estimator API正在积极开发中,在tensorflow的主分支中,input_fn已经可以返回数据集了,所以也许我要求得太早,这个功能还没有准备好.但如果是这样,请提供可以跟踪此实施的票据.

Answer 1

使用tf.data.Dataset.cache()确实不是一个好选择,因为它会将整个数据集缓存到内存中,这需要时间并且可能会溢出内存.

tf.data.Dataset.prefetch()要做的就是在管道的末尾使用,这将始终确保数据管道包含buffer_size元素.最后通常就足够buffer_size = 1了:

dataset = ...
dataset = dataset.batch(128)
dataset = dataset.prefetch(1)  # prefetch one batch

正如@mrry在本回答中所解释的那样,您也可以尝试增加预取批次的数量.

通常,在管道的最末端添加一个小的预取缓冲区(可能只有一个元素)是最有用的,但是更复杂的管道可以从额外的预取中受益,特别是当生成单个元素的时间可能变化时.

如果与GPU计算相比仍然有慢速输入管道,则需要使用num_parallel_calls参数来增加并行工作的线程数tf.data.Dataset.map().