From a0ad3bd283e8263b89cb16d977a6931c11cb0499 Mon Sep 17 00:00:00 2001
From: Michael <m.langguth@fz-juelich.de>
Date: Thu, 7 Apr 2022 15:48:11 +0200
Subject: [PATCH] Try caching dataset.
---
.../model_modules/video_prediction/datasets/era5_dataset.py | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/video_prediction_tools/model_modules/video_prediction/datasets/era5_dataset.py b/video_prediction_tools/model_modules/video_prediction/datasets/era5_dataset.py
index 79e5e2fd..f623636a 100644
--- a/video_prediction_tools/model_modules/video_prediction/datasets/era5_dataset.py
+++ b/video_prediction_tools/model_modules/video_prediction/datasets/era5_dataset.py
@@ -180,7 +180,9 @@ class ERA5Dataset(object):
if shuffle:
random.shuffle(filenames)
dataset = tf.data.TFRecordDataset(filenames, buffer_size = 8* 1024 * 1024)
-
+
+ # cache dataset
+ dataset = dataset.cache()
# dataset = dataset.filter(self.filter)
if shuffle:
dataset = dataset.apply(tf.contrib.data.shuffle_and_repeat(buffer_size =1024, count = self.num_epochs))
--
GitLab