TensorFlow 数据集中没有写入缓存文件

Question

我正在尝试管理一个内存无法容纳的大型图像数据集，同时需要进行一些特定的计算。目前，我的代码如下所示：

  files  = [str(f) for f in self.files]
  labels = self.categories
  batch_size= 32


  dataset = tf.data.Dataset.from_generator(
        lambda: zip(files, labels),
        output_types=(tf.string, tf.uint8),
        output_shapes=(tf.TensorShape([]), tf.TensorShape([]))
  )

  dataset = dataset.map(
        lambda x, y: tf.py_function(_parser, [x, y, category_count], [tf.float32, tf.uint8]),
        num_parallel_calls=tf.data.experimental.AUTOTUNE,
        deterministic=False)

  dataset.cache(filename='/tmp/dataset.tmp')

  if mode == tf.estimator.ModeKeys.TRAIN:
        dataset = dataset.shuffle(buffer_size=10*batch_size, reshuffle_each_iteration=True)

  dataset = dataset.batch(batch_size=batch_size, drop_remainder=False)

  if mode == tf.estimator.ModeKeys.TRAIN:
        dataset.repeat(None)
  else:
        dataset.repeat(1)

  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

_parser() 函数打开一个图像文件，进行一系列变换，returns 一个张量和一个 1-hot 编码向量。缓存步骤似乎无法正常工作，但是：

第一个epoch和后面的epoch之间的计算时间没有显着改善
虽然交换分区几乎已满 (~90%)，但在此过程中没有创建缓存文件

cache()函数是否只有在内存和交换分区都满时才创建文件？此外，我希望一次只读取 batch_size 个文件。但是，似乎所有文件都在映射步骤中一次读取。我应该考虑结合使用 interleave() 和 from_generator() 吗？或者我应该先批处理文件，然后映射它们？

Answer 1

注意当数据集较小时应该使用cache()。如果数据集很大（在您的情况下），RAM 将不足以缓存其内容，因此它不适合内存。您应该增加 RAM 的容量或采用其他方法来加速训练。

训练速度变慢的另一个原因是使用 map() 函数时的预处理阶段。
map() 方法将转换应用于每个项目，这与 apply() 方法将转换应用于整个数据集不同。

您可以使用 interleave() 并保留 map() 和 batch() 的相同顺序。
您已经通过设置 num_parallel_calls 并将其设置为 tf.data.experimental.AUTOTUNE 来使用线程，充分利用可用的资源。

你也可以规范化你的输入数据，然后cache，如果它再次不适合内存那么最好不要在大型数据集上cache。

您可以遵循来自 TensorFlow 的 these 性能提示。

如果你有多个 workers/devices 它将帮助你加快训练速度。

下面是显示多线程加载和预处理预取的示例图。

TensorFlow 数据集中没有写入缓存文件

No cache file written in TensorFlow dataset

tensorflow-datasets

tensorflow2.0