Repository: OpenNMT/OpenNMT-tf
Pull request #645 (2020-04-04 08:37:21)

Add parameter max_exports_to_keep to limit number of saved exports


Patch:
diff --git a/docs/configuration.md b/docs/configuration.md
index 48c75243..521df76d 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -289,6 +289,8 @@ eval:
   export_on_best: bleu
   # (optional) Format of the exported model (can be: "saved_model, "ctranslate2", default: "saved_model").
   export_format: saved_model
+  # (optional) Maximum number of exports to keep on disk (default: 5).
+  max_exports_to_keep: 5
 
   # (optional) Early stopping condition.
   # Should be read as: stop the training if "metric" did not improve more
diff --git a/opennmt/evaluation.py b/opennmt/evaluation.py
index 2a673d0d..6a7bd378 100644
--- a/opennmt/evaluation.py
+++ b/opennmt/evaluation.py
@@ -2,6 +2,7 @@
 
 import collections
 import os
+import shutil
 
 import tensorflow as tf
 
@@ -32,7 +33,8 @@ class Evaluator(object):
                early_stopping=None,
                model_dir=None,
                export_on_best=None,
-               exporter=None):
+               exporter=None,
+               max_exports_to_keep=5):
     """Initializes the evaluator.
 
     Args:
@@ -50,6 +52,8 @@ class Evaluator(object):
         best value so far.
       exporter: A :class:`opennmt.utils.Exporter` instance to export the model.
         Defaults to :class:`opennmt.utils.SavedModelExporter`.
+      max_exports_to_keep: Maximum number of exports to keep. Older exports will
+        be garbage collected. Set to ``None`` to keep all exports.
 
     Raises:
       ValueError: If :obj:`save_predictions` is set but the model is not compatible.
@@ -119,6 +123,7 @@ class Evaluator(object):
     self._export_on_best = export_on_best
     self._exporter = exporter
     self._export_dir = export_dir
+    self._max_exports_to_keep = max_exports_to_keep
 
   @classmethod
   def from_config(cls, model, config, features_file=None, labels_file=None):
@@ -163,7 +168,8 @@ class Evaluator(object):
         early_stopping=early_stopping,
         model_dir=config["model_dir"],
         export_on_best=eval_config.get("export_on_best"),
-        exporter=exporters.make_exporter(eval_config.get("export_format", "saved_model")))
+        exporter=exporters.make_exporter(eval_config.get("export_format", "saved_model")),
+        max_exports_to_keep=eval_config.get("max_exports_to_keep", 5))
 
   @property
   def predictions_dir(self):
@@ -302,6 +308,7 @@ class Evaluator(object):
 
     self._record_results(step, results)
     self._maybe_export(step, results)
+    self._maybe_garbage_collect_exports()
     return results
 
   def _record_results(self, step, results):
@@ -326,6 +333,16 @@ class Evaluator(object):
                          export_dir, self._export_on_best, results[self._export_on_best])
     self._model.export(export_dir, exporter=self._exporter)
 
+  def _maybe_garbage_collect_exports(self):
+    if self._max_exports_to_keep is None or not os.path.exists(self._export_dir):
+      return
+    exported_steps = list(sorted(map(int, os.listdir(self._export_dir))))
+    num_exports = len(exported_steps)
+    if num_exports > self._max_exports_to_keep:
+      steps_to_remove = exported_steps[:num_exports - self._max_exports_to_keep]
+      for step in steps_to_remove:
+        shutil.rmtree(os.path.join(self._export_dir, str(step)))
+
 
 def early_stop(metrics, steps, min_improvement=0, higher_is_better=False):
   """Early stopping condition.


Test patch:
diff --git a/opennmt/tests/evaluation_test.py b/opennmt/tests/evaluation_test.py
index 0721eaef..0f9350b9 100644
--- a/opennmt/tests/evaluation_test.py
+++ b/opennmt/tests/evaluation_test.py
@@ -70,7 +70,6 @@ class EvaluationTest(tf.test.TestCase):
       features.write("1\n2\n")
       labels.write("1\n2\n")
     model = TestModel({"a": [2, 5, 8], "b": [3, 6, 9]})
-    model.initialize({})
     early_stopping = evaluation.EarlyStopping(metric="loss", min_improvement=0, steps=1)
     evaluator = evaluation.Evaluator(
         model,
@@ -130,7 +129,6 @@ class EvaluationTest(tf.test.TestCase):
       features.write("1\n2\n")
       labels.write("1\n2\n")
     model = TestModel()
-    model.initialize({})
     evaluator = evaluation.Evaluator(
         model,
         features_file,
@@ -143,6 +141,38 @@ class EvaluationTest(tf.test.TestCase):
     self.assertIn("rouge-2", evaluator.metrics_name)
     self.assertIn("rouge-l", evaluator.metrics_name)
 
+  def testExportsGarbageCollection(self):
+    features_file = os.path.join(self.get_temp_dir(), "features.txt")
+    labels_file = os.path.join(self.get_temp_dir(), "labels.txt")
+    model_dir = self.get_temp_dir()
+    with open(features_file, "w") as features, open(labels_file, "w") as labels:
+      features.write("1\n2\n")
+      labels.write("1\n2\n")
+    model = TestModel()
+    exporter = TestExporter()
+    evaluator = evaluation.Evaluator(
+        model,
+        features_file,
+        labels_file,
+        batch_size=1,
+        model_dir=model_dir,
+        export_on_best="loss",
+        exporter=exporter,
+        max_exports_to_keep=2)
+
+    # Generate some pre-existing exports.
+    for step in (5, 10, 15):
+      exporter.export(model, os.path.join(evaluator.export_dir, str(step)))
+
+    def _eval_step(step, loss, expected_exported_steps):
+      model.next_loss.assign(loss)
+      evaluator(step)
+      exported_steps = list(sorted(map(int, os.listdir(evaluator.export_dir))))
+      self.assertListEqual(exported_steps, expected_exported_steps)
+
+    _eval_step(20, 3, [15, 20])  # Exports 5 and 10 should be removed.
+    _eval_step(25, 2, [20, 25])  # Export 15 should be removed.
+
   def testEarlyStop(self):
     self.assertFalse(
         evaluation.early_stop([3.1, 2.7, 2.6], 4))