diff --git a/modules/deeplearning/cloud_opd_fcn_abi.py b/modules/deeplearning/cloud_opd_fcn_abi.py index fa1a5c6e3d0a0c5eba69898edbcaa59f686e4109..5b14e79b861b4d360c3b8d9fb2dd302fa23d80af 100644 --- a/modules/deeplearning/cloud_opd_fcn_abi.py +++ b/modules/deeplearning/cloud_opd_fcn_abi.py @@ -377,7 +377,7 @@ class SRCNN: dataset = tf.data.Dataset.from_generator(gen, output_types=tf.int32) dataset = dataset.batch(PROC_BATCH_SIZE) dataset = dataset.map(self.data_function, num_parallel_calls=8) - dataset = dataset.cache() + dataset = dataset.cache(CACHE_FILE) dataset = dataset.shuffle(PROC_BATCH_BUFFER_SIZE, reshuffle_each_iteration=True) if DO_AUGMENT: dataset = dataset.map(augment_image(), num_parallel_calls=8) @@ -396,7 +396,7 @@ class SRCNN: dataset = tf.data.Dataset.from_generator(gen, output_types=tf.int32) dataset = dataset.batch(PROC_BATCH_SIZE) dataset = dataset.map(self.data_function_test, num_parallel_calls=8) - dataset = dataset.cache() + dataset = dataset.cache(CACHE_FILE) self.test_dataset = dataset def setup_pipeline(self, train_data_files, train_label_files, test_data_files, test_label_files, num_train_samples):