使用比 Ram 更多数据的 TensorFlow 数据集和估计器

Question

我最近将我的建模框架切换为使用自定义 Tensorflow Estimators 和 Datasets，总体上对这个工作流程非常满意。

但是，我刚刚注意到我的 dataset_input_fn 如何从 tfrecords 加载数据的问题。我的输入函数是根据 Tensorflow 文档中的 example 建模的。当我的示例超出 RAM 容量时，就会出现问题。如果我有 1e6 个示例，并将我的 shuffle buffer_size 设置为 1e5，则选择一次 1e5 示例的子集，打乱，然后迭代。这意味着我的模型只在我整个数据集的 10% 上进行了训练。我设置此行为的代码完全是从Tensorflow documentation example code 借来的：

dataset = dataset.map(parser)
dataset = dataset.shuffle(buffer_size=10000)
dataset = dataset.batch(32)
dataset = dataset.repeat(num_epochs)
iterator = dataset.make_one_shot_iterator()

我的问题：是否可以在我训练时用初始 1e5 之外的新示例填充随机播放缓冲区？ one_shot_iterator 是否支持这种类型的功能？我需要使用可初始化的迭代器吗？

谢谢！

Answer 1

我发现目前看来可行的解决方法。通过一些实验，我了解到在实例化 TFRecordDataset 时，

filenames = ["file1.tfrecord", ..., "filen.tfrecord"]
dataset = tf.data.TFRecordDataset(filenames)

并设置洗牌缓冲区：

 dataset = dataset.shuffle(buffer_size=10000)

缓冲区仅填充了前 10000 个示例，无论需要多少 tf 记录。例如，就我而言，我有大约 300 个 tfrecord 文件，每个文件包含 4096 个示例。经过检查，我的 shuffle 缓冲区似乎只包含我的文件名列表中前 3 条 tf 记录的示例。由于我的文件名列表是静态的，这意味着我的模型只训练了我的前 3 个 tfrecord！

我现在的解决方法很简单。在我的训练循环中，我已经在 Estimator.train 和 Estimator.evaluate 之间交替，我注意到每次调用 Estimator.train 时，都会重新填充随机缓冲区。然后，我的解决方案是每次调用 input_fn 时都洗牌我的文件名。这不是一个非常优雅的解决方案，但确实达到了允许 my 遍历所有 tf​​record 的预期效果。

#My Crappy Fix: shuffle file names in input_fn
np.random.shuffle(filenames)
dataset = tf.data.TFRecordDataset(filenames)

这个解决方案令人讨厌的地方（除了它的笨拙之外）是我的小批量不是“全局随机的”。相反，它们是从 tf 记录的一小部分中选择的，并且只有该子集用于每个训练/评估周期。缓解这种情况的一种方法是增加我的 shuffle 缓冲区大小或减少我的 tfrecord 大小，我可能会同时做这两个。最后，我认为值得注意的是，如果

shuffle_buffer_size < (tf_record_size + minibatch_size) 

然后，据我所知，我的 TFRecordDataset 将从单个 tfrecord 文件中提取！

最后，我认为relevant tensorflow documentation 不能很好地传达这些复杂性。该文档提到了在不适合内存的大型数据集上进行训练的能力，但没有提供太多细节。 tf 作者在写这篇文章时似乎不太可能想到我的 hacky 策略，所以我仍然很好奇是否有更好的方法。