avoid iterating dataset when window it.

Pandinosaurus · Sep 30, 2020 · 59128d3 · 59128d3
1 parent 4f4d2a4
commit 59128d3
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/autokeras/preprocessors/common.py b/autokeras/preprocessors/common.py
@@ -84,11 +84,9 @@ def __init__(self, lookback, batch_size, **kwargs):
     def transform(self, dataset):
         dataset = dataset.unbatch()
         dataset = dataset.window(self.lookback, shift=1, drop_remainder=True)
-        final_data = []
-        # TODO: Avoid iterating the dataset to speedup and save memory.
-        for window in dataset:
-            final_data.append([elems.numpy() for elems in window])
-        dataset = tf.data.Dataset.from_tensor_slices(final_data)
+        dataset = dataset.flat_map(
+            lambda x: x.batch(self.lookback, drop_remainder=True)
+        )
         dataset = dataset.batch(self.batch_size)
         return dataset