tensorflow - Tensorflow 数据集 API 中的内存管理

Question

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

Answer 1

是的。官方指南中的示例(使用数据集 API 用于 TensorFlow 输入管道，https://www.tensorflow.org/programmers_guide/datasets)

filenames = ["/var/data/file1.tfrecord", "/var/data/file2.tfrecord"]
dataset = tf.contrib.data.TFRecordDataset(filenames)
dataset = dataset.map(...) ## Parsing data with a user specified function
dataset = dataset.shuffle(buffer_size=10000) ## 10000: size of sample/record pool for random selection
dataset = dataset.batch(32) ## 32: number of samples/records per batch (to be read into memory)
dataset = dataset.repeat() ## None: keep repeating