Start respecting keep_checkpoint_max

Change: 125712724

tensorflow/contrib/learn/python/learn/estimators/estimator.py

@@ -496,6 +496,7 @@ class BaseEstimator(sklearn.BaseEstimator):
           supervisor_is_chief=is_chief,
           supervisor_master=self._config.master,
           supervisor_save_model_secs=self._config.save_checkpoints_secs,
+          keep_checkpoint_max=self._config.keep_checkpoint_max,
           feed_fn=feed_fn,
           steps=steps,
           fail_on_nan_loss=fail_on_nan_loss,

tensorflow/contrib/learn/python/learn/graph_actions.py

@@ -87,10 +87,12 @@ class NanLossDuringTrainingError(RuntimeError):
     return 'NaN loss during training.'
 
 
-def _make_saver(graph):
+def _make_saver(graph, keep_checkpoint_max=5):
   vars_to_save = graph.get_collection(ops.GraphKeys.VARIABLES)
   if vars_to_save:
-    return tf_saver.Saver(vars_to_save, sharded=True)
+    return tf_saver.Saver(vars_to_save,
+                          sharded=True,
+                          max_to_keep=keep_checkpoint_max)
   else:
     return None
 
@@ -136,6 +138,7 @@ def train(graph,
           supervisor_is_chief=True,
           supervisor_master='',
           supervisor_save_model_secs=600,
+          keep_checkpoint_max=5,
           supervisor_save_summaries_steps=100,
           feed_fn=None,
           steps=None,
@@ -176,6 +179,10 @@ def train(graph,
     supervisor_master: The master string to use when preparing the session.
     supervisor_save_model_secs: Save a checkpoint every
       `supervisor_save_model_secs` seconds when training.
+    keep_checkpoint_max: The maximum number of recent checkpoint files to
+      keep. As new files are created, older files are deleted. If None or 0,
+      all checkpoint files are kept. This is simply passed as the max_to_keep
+      arg to tf.Saver constructor.
     supervisor_save_summaries_steps: Save summaries every
       `supervisor_save_summaries_steps` seconds when training.
     feed_fn: A function that is called every iteration to produce a `feed_dict`
@@ -241,7 +248,7 @@ def train(graph,
       init_feed_dict=init_feed_dict,
       is_chief=supervisor_is_chief,
       logdir=output_dir,
-      saver=_make_saver(graph),
+      saver=_make_saver(graph, keep_checkpoint_max),
       global_step=global_step_tensor,
       summary_op=None,
       summary_writer=summary_writer,