tensorflow - 有没有办法在内存中保存一个 Tensorflow 记录文件？

Question

情况是这样的:我正在处理一个大型 Tensorflow 记录文件。大约 50 GB。但是，我正在执行此培训的机器有 128 GB 的 RAM。 50 小于 128，因此即使这是一个大文件，您也会认为可以将其保存在内存中并节省慢速 I/O 操作符。但我使用的是 TFRecordDataset 类，似乎整个 TFRecord 系统都是专门为不这样做而设计的，而且我看不出有任何方法可以强制它在内存中保存记录。由于它在每个时期都会重新加载它们，因此我在从 50 GB 文件读取缓慢的 I/O 操作上浪费了过多的时间。

我想我可以在 python 中将记录加载到内存中，然后使用 feed_dict 将它们一个一个地加载到我的模型中，绕过整个数据集类。但这似乎是一种不太优雅的处理方式，需要进行一些重新设计。如果我可以强制 TFRecordDataset 将所有内容加载到内存中并在各个时期之间将其保存在那里，一切都会简单得多......

Answer 1

你需要tf.data.Dataset.cache()手术。要达到预期的效果(将文件保存在内存中)，请将其放在 TFRecordDataset 之后并且不向其提供任何参数:

  dataset = tf.data.TFRecordDataset(filenames)
  dataset = dataset.cache()

当 cache() 操作在没有参数的情况下被调用时，缓存是在内存中完成的。

此外，如果您对这些记录进行了一些后处​​理，例如 dataset.map(...)，那么将 cache() 在输入管道的末端运行。

更多信息可以在“输入管道性能指南”中找到Map and Cache section .