Extending GeneratorEnqueuer to handle finite generators.

keras/utils/data_utils.py

@@ -537,10 +537,13 @@ def stop(self, timeout=None):
 class GeneratorEnqueuer(SequenceEnqueuer):
     """Builds a queue out of a data generator.
 
+    The provided generator can be finite in which case the class will throw
+    a `StopIteration` exception.
+
     Used in `fit_generator`, `evaluate_generator`, `predict_generator`.
 
     # Arguments
-        generator: a generator function which endlessly yields data
+        generator: a generator function which yields data
         use_multiprocessing: use multiprocessing if True, otherwise threading
         wait_time: time to sleep in-between calls to `put()`
         random_seed: Initial seed for workers,
@@ -576,6 +579,8 @@ def data_generator_task():
                         self.queue.put(generator_output)
                     else:
                         time.sleep(self.wait_time)
+                except StopIteration:
+                    break
                 except Exception:
                     self._stop_event.set()
                     raise
@@ -648,4 +653,8 @@ def get(self):
                 if inputs is not None:
                     yield inputs
             else:
-                time.sleep(self.wait_time)
+                all_finished = all([not thread.is_alive() for thread in self._threads])
+                if all_finished:
+                    raise StopIteration()
+                else:
+                    time.sleep(self.wait_time)

tests/keras/utils/data_utils_test.py

@@ -245,5 +245,40 @@ def test_ordered_enqueuer_fail_processes():
         next(gen_output)
 
 
+@threadsafe_generator
+def create_finite_generator_from_sequence_threads(ds):
+    for i in range(len(ds)):
+        yield ds[i]
+
+
+def create_finite_generator_from_sequence_pcs(ds):
+    for i in range(len(ds)):
+        yield ds[i]
+
+
+def test_finite_generator_enqueuer_threads():
+    enqueuer = GeneratorEnqueuer(create_finite_generator_from_sequence_threads(
+        TestSequence([3, 200, 200, 3])), use_multiprocessing=False)
+    enqueuer.start(3, 10)
+    gen_output = enqueuer.get()
+    acc = []
+    for output in gen_output:
+        acc.append(int(output[0, 0, 0, 0]))
+    assert len(set(acc) - set(range(100))) == 0, "Output is not the same"
+    enqueuer.stop()
+
+
+def test_finite_generator_enqueuer_processes():
+    enqueuer = GeneratorEnqueuer(create_finite_generator_from_sequence_pcs(
+        TestSequence([3, 200, 200, 3])), use_multiprocessing=True)
+    enqueuer.start(3, 10)
+    gen_output = enqueuer.get()
+    acc = []
+    for output in gen_output:
+        acc.append(int(output[0, 0, 0, 0]))
+    assert acc != list(range(100)), "Order was keep in GeneratorEnqueuer with processes"
+    enqueuer.stop()
+
+
 if __name__ == '__main__':
     pytest.main([__file__])