Есть ли способ сохранить файл записи Tensorflow в памяти?

Question

Вот ситуация: я работаю с большим файлом записей Tensorflow. Это около 50 ГБ. Однако машина, на которой я тренируюсь, имеет 128 ГБ ОЗУ. 50 меньше 128, поэтому, даже если это большой файл, можно подумать, что его можно будет сохранить в памяти и сэкономить на медленных операторах ввода-вывода. Но я использую класс TFRecordDataset, и похоже, что вся система TFRecord разработана специально для того, чтобы этого не делать, и я не вижу способа заставить ее хранить записи в памяти. И так как они перезагружают их каждую эпоху, я трачу непомерное количество времени на медленные операции ввода-вывода, читая этот файл размером 50 ГБ.

Я полагаю, что я мог бы загрузить записи в память на python, а затем загрузить их в мою модель одну за другой с помощью feed_dict, минуя весь класс Dataset. Но это кажется менее элегантным способом управления вещами и потребует некоторого изменения дизайна. Все было бы намного проще, если бы я мог просто заставить TFRecordDataset загружать все в память и сохранять ее там между эпохами...

Answer 1

Вам нужна операция tf.data.Dataset.cache(). Чтобы добиться желаемого эффекта (сохранение файла в памяти), поместите его сразу после TFRecordDataset и не предоставляйте ему никаких аргументов:

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

Когда cache() операция вызывается без аргументов, чем кеширование в памяти.

Также, если у вас есть постобработка этих записей, например, с dataset.map(...)тогда было бы еще выгоднее поставить cache() операция в конце входного конвейера.

Дополнительную информацию можно найти в разделе " Карта и кэш" в "Руководстве по производительности входного конвейера".