首页 > 解决方案 > 仅使用 tf.data API(无 feed_dict)的高性能输入管道的最佳实践

问题描述

官方的 TensorFlow性能指南规定如下:

虽然使用 feed_dict 提供数据提供了高度的灵活性,但通常 feed_dict 不提供可扩展的解决方案。如果只使用单个 GPU,tf.data API 和 feed_dict 性能之间的差异可能可以忽略不计。我们的建议是避免将 feed_dict 用于除琐碎示例之外的所有示例。特别是,避免将 feed_dict 与大输入一起使用。

但是,完全避免使用feed_dict似乎是不可能的。考虑以下带有训练、验证和测试数据集的设置。

ds = tf.data.Dataset
n_files = 1000 # total number of tfrecord files
split = int(.67 * n_files)
files = ds.zip((ds.range(n_files),ds.list_files("train/part-r-*")))
train_files = files.filter(lambda a, b: a < split).map(lambda a,b: b)
validation_files = files.filter(lambda a, b: a >= split).map(lambda a,b: b)
test_files = ds.list_files("test/part-r-*")

解析数据集的常用方法可能如下所示:

def setup_dataset(self, file_ds, mode="train"):

   data = file_ds.apply(tf.contrib.data.parallel_interleave(
       tf.data.TFRecordDataset,
       cycle_length=4,
       sloppy=True,
       buffer_output_elements=self.batch_size * 8,
       prefetch_input_elements=self.batch_size * 8
   ))

   if mode == "train":
       data = data.map(self.train_data_parser)
   else:
       data = data.map(self.test_data_parser)

   return data

然后,不是通过feed_dict in提供各个功能,而是使用orsession.run()创建一个可重用的迭代器。我将展示一个与前者的例子,但无论哪种方式你都会遇到同样的问题。Iterator.from_structure()Iterator.from_string_handle()

train = self.setup_dataset(train_files)
self.ops["template_iterator"] = tf.data.Iterator.from_structure(train.output_types, train.output_shapes)
self.ops["next_batch"] = self.ops["template_iterator"].get_next(name="next_batch")
self.ops["train_init"] = self.ops["template_iterator"].make_initializer(train)

validation = self.setup_dataset(validation_files)
self.ops["validation_init"] = self.ops["template_iterator"].make_initializer(validation)

这一切都很好,但是我应该如何处理测试数据集?测试数据集将不包含标签特征,因此与训练和验证数据集不符合相同的output_typesoutput_shapes 。

理想情况下,我希望从 SavedModel 恢复并初始化测试数据集,而不是通过 API 提供模型。

在推理过程中我缺少什么技巧来合并测试数据集?

标签: python-3.xtensorflow

解决方案


我为训练和推理设置了数据集和迭代器,如下所示:

# Train dataset
images_train = tf.placeholder(tf.float32, train_images.shape)
labels_train = tf.placeholder(tf.float32, train_masks.shape)
dataset_train = tf.data.Dataset.from_tensor_slices({"images": images_train, "masks": labels_train})
dataset_train = dataset_train.batch(MINIBATCH)
dataset_train = dataset_train.map(lambda x: map_helper(x, augmentation), num_parallel_calls=8)
dataset_train = dataset_train.shuffle(buffer_size=10000)

iterator_train = tf.data.Iterator.from_structure(dataset_train.output_types, dataset_train.output_shapes)
training_init_op = iterator_train.make_initializer(dataset_train)
batch_train = iterator_train.get_next()

# Inference dataset
images_infer = tf.placeholder(tf.float32, shape=[None] + list(valid_images.shape[1:]))
labels_infer = tf.placeholder(tf.float32, shape=[None] + list(valid_masks.shape[1:]))
dataset_infer = tf.data.Dataset.from_tensor_slices({"images": images_infer, "masks": labels_infer})
dataset_infer = dataset_infer.batch(MINIBATCH)

iterator_infer = tf.data.Iterator.from_structure(dataset_infer.output_types, dataset_infer.output_shapes)
infer_init_op = iterator_infer.make_initializer(dataset_infer)
batch_infer = iterator_infer.get_next()

训练

初始化用于训练的迭代器training_init_op

sess.run(training_init_op, feed_dict={images_train: train_images, labels_train: train_masks})

验证

使用初始化推理迭代器进行验证infer_init_op

sess.run(infer_init_op, feed_dict={images_infer: images_val, labels_infer: masks_val})

测试

使用 初始化推理迭代器以进行测试infer_init_op。这有点 hacky,但是我创建了一个带有零的数组,其中标签将使用与验证相同的迭代器

sess.run(infer_init_op, feed_dict={images_infer: images_test, labels_infer: np.zeros(images_test.shape)})

或者,您可以为训练/验证/测试创建 3 个不同的数据集/迭代器


推荐阅读