如何衡量 Tensorflow 新 DataSet API 中 Map 函数的性能？

Posted 2023-03-27

【中文标题】如何衡量 Tensorflow 新 DataSet API 中 Map 函数的性能？

【英文标题】：How To Measure The Performance Of Map Function In Tensorflow's New DataSet API?

【发布时间】：2019-06-13 12:34:45

【问题描述】：

我一直在我的 GPU 实例中使用 Tensorflow 1.12 版，我有大约 130 个 TfRecords 文件，其中包含 120 万个 ImageNet 数据。首先我应用了一个地图函数，然后flat_map 来扩充数据集，最终将是 120 万 x 2048 图像。

self.filenames = tf.placeholder(tf.string, shape=[None])
self.eval_filenames = tf.placeholder(tf.string, shape = [None])
dataset = tf.data.TFRecordDataset(self.filenames)
eval_dataset = tf.data.TFRecordDataset(self.eval_filenames)
print("inside dataset ", dataset.output_shapes)
dataset = dataset.map(self.decode, num_parallel_calls=10)
dataset = dataset.flat_map(self.apply_flip_crop)
dataset = dataset.batch(self.config["batch_size"])
dataset = dataset.prefetch(2)
iterator = dataset.make_initializable_iterator()

这里的 decode 函数返回图像的扁平化数组和 one-hot 编码标签。然而，flat_map 中传递的函数做了一件很繁重的事情，就像：两个循环来创建切片并反转它们，每个循环产生 1024 个张量。单个图像的最终输出将是 [2048, 224, 224, 3] 张量。函数如下所示：

   def apply_flip_crop(self, tf_example, lable):
        """
        Calls a helper function random_crop flips which randomly crops and flips
        the images, and returns the agumented tensors.

        Parameters
        ----------
        :param tf_example: A tensor of shape [batchsize, flattedimageshape]
        :type tf_example: Tensors [batchsize, flattedimageshape]
        :param lable: A Constant integer representing the class_id of the image.
        :type lable: tf.int32
        :return: Tensors of shape [flattedimageshape], label of image tf.int32
        :rtype: Tensors

        """

        data = tf.reshape(tf_example, [256, 256, 3])
        data = self.random_crop_flip(data)
        lables = [lable for i in range(2048)]
        return tf.data.Dataset.from_tensor_slices((data, lables))


    def random_crop_flip(self, image):
        """
        Apply random crop and random flip to the image tensor.

        Parameters
        ----------
        :param image: A tensor representing a flattened image array.
        :type image: Tensor of shape [imageflattenedarray]
        :return: List of 2048 tensors of shape [imageflattenedarray]
        :rtype: List

        """
        crops = []

        for i in range(256 - 224):
            for j in range(256 - 224):
                crop = tf.slice(image, [i, j, 0], [224, 224, 3])
                crop2 = tf.reverse(crop, axis=[1])
                crops.append(crop)
                crops.append(crop2)
        return crops

现在的问题是训练过程非常缓慢。我读过dataset.from_tensor_slices 对这种需求非常不利。但我认为这里有很多可以改进的地方。为此，我需要可视化每个操作的性能。主要是flat_map函数。 我正在使用 tensorflow 的 RunTime Statistic，如下所示：

sess.run(iterator.initializer, feed_dict=data_gen.filenames:
                                                      training_filenames,
                 options=run_options, run_metadata=run_metadata)
next_element = iterator.get_next()
for i in range(1):

     datapoint = sess.run(next_element, options=run_options, 
                run_metadata=run_metadata)

      summary_writer.add_run_metadata(run_metadata, 'step%d' % i)

它记录了准备数据集所花费的时间，但它没有记录执行flat_map 操作所花费的时间，这是我怀疑的问题，这是性能滞后的地方。

感谢您在性能建议以及flat_map函数所用时间的测量方面的帮助。

提前致谢。

【参考方案1】：

尝试使用 TensorFlow 的时间线：HowTo profile TensorFlow