Tensorflow数据集API中的内存管理

Question

我的训练数据集太大而无法容纳到内存中,因此我的代码一次只能从磁盘读取1,000条记录.现在我想使用Tensorflow的新数据集API.数据集API是否允许我指定要保留在内存中的记录数,或者Tensorflow是否自动管理内存以便我不必？

Answer 1

如果您将通过batch_size指定记录数。在这种情况下，TF 将仅从文件中获取batch_size 元素。您还可以指定shuffle，这将保证内存中的所有时间都处于最大buffer_size元素数。

我在我的 tfrecords 文件上验证了它。我有 100 个 tfrecords 文件，每个文件大约 10Gb（这比我笔记本电脑上的内存还要大）。一切都很好。