Change StopAfterNEvalsHook to support iterating dataset indefinitely.

PiperOrigin-RevId: 169415251

1 files changed, 7 insertions, 3 deletions

diff --git a/tensorflow/python/training/evaluation.py b/tensorflow/python/training/evaluation.py
index 3baf1541aa..fdcb9c2e90 100644
--- a/tensorflow/python/training/evaluation.py
+++ b/tensorflow/python/training/evaluation.py
@@ -83,7 +83,8 @@ class _StopAfterNEvalsHook(session_run_hook.SessionRunHook):
     """Constructs the run hook.
 
     Args:
-      num_evals: The number of evaluations to run for.
+      num_evals: The number of evaluations to run for. if set to None, will
+        iterate the dataset until all inputs are exhausted.
       log_progress: Whether to log evaluation progress, defaults to True.
     """
     # The number of evals to run for.
@@ -102,8 +103,11 @@ class _StopAfterNEvalsHook(session_run_hook.SessionRunHook):
   def after_run(self, run_context, run_values):
     evals_completed = run_values.results['evals_completed']
     if self._log_progress:
-      logging.info('Evaluation [%d/%d]', evals_completed, self._num_evals)
-    if evals_completed >= self._num_evals:
+      if self._num_evals is None:
+        logging.info('Evaluation [%d]', evals_completed)
+      else:
+        logging.info('Evaluation [%d/%d]', evals_completed, self._num_evals)
+    if self._num_evals is not None and evals_completed >= self._num_evals:
       run_context.request_stop()