Repository: OpenNMT/OpenNMT-tf
Pull request #818 (2021-03-23 09:19:58)

Forward model description file to generated checkpoint directories


Patch:
diff --git a/opennmt/config.py b/opennmt/config.py
index 277a11d7..6df91b04 100644
--- a/opennmt/config.py
+++ b/opennmt/config.py
@@ -14,6 +14,9 @@ from opennmt.schedules import lr_schedules as schedules_lib
 from opennmt.utils.misc import merge_dict
 
 
+MODEL_DESCRIPTION_FILENAME = "model_description.py"
+
+
 def load_model_module(path):
     """Loads a model configuration file.
 
@@ -102,7 +105,7 @@ def load_model(
     """
     if model_file and model_name:
         raise ValueError("only one of model_file and model_name should be set")
-    model_description_path = os.path.join(model_dir, "model_description.py")
+    model_description_path = os.path.join(model_dir, MODEL_DESCRIPTION_FILENAME)
 
     if model_file:
         model = load_model_from_file(model_file, as_builder=as_builder)
diff --git a/opennmt/runner.py b/opennmt/runner.py
index ec7464dc..1aa9b474 100644
--- a/opennmt/runner.py
+++ b/opennmt/runner.py
@@ -5,6 +5,7 @@ import os
 import sys
 import random
 import math
+import shutil
 import subprocess
 import tempfile
 import yaml
@@ -16,6 +17,7 @@ from opennmt import evaluation
 from opennmt import inference
 from opennmt import models
 from opennmt import training as training_util
+from opennmt.config import MODEL_DESCRIPTION_FILENAME
 from opennmt.utils import checkpoint as checkpoint_util
 from opennmt.utils import misc
 from opennmt.version import __version__
@@ -343,6 +345,7 @@ class Runner(object):
         output_dir = checkpoint_util.average_checkpoints(
             checkpoint.model_dir, output_dir, trackables, max_count=max_count
         )
+        _forward_model_description(self.model_dir, output_dir)
         self._config["model_dir"] = output_dir
         return output_dir
 
@@ -372,6 +375,7 @@ class Runner(object):
         )
         cur_checkpoint.restore()
         model.create_variables(optimizer=optimizer)
+        source_dir = self.model_dir
 
         self._config["model_dir"] = output_dir
         if src_vocab is not None:
@@ -391,6 +395,7 @@ class Runner(object):
         )
         new_optimizer.iterations.assign(optimizer.iterations)
         new_checkpoint.save()
+        _forward_model_description(source_dir, output_dir)
         return output_dir
 
     def infer(
@@ -473,6 +478,15 @@ class Runner(object):
         )
 
 
+def _forward_model_description(source, destination):
+    source = os.path.join(source, MODEL_DESCRIPTION_FILENAME)
+    if os.path.isfile(source):
+        if not os.path.isdir(destination):
+            os.makedirs(destination)
+        destination = os.path.join(destination, MODEL_DESCRIPTION_FILENAME)
+        shutil.copyfile(source, destination)
+
+
 def _count_batch_accum(batch_size, target_batch_size, num_replicas=1):
     """Given the current batch size, the number of replicas, and the requested
     effective batch size, returns the number of gradients to accumulate.
@@ -522,7 +536,7 @@ def _auto_tune_batch_size(
         min_range,
     )
 
-    model_description = os.path.join(config["model_dir"], "model_description.py")
+    model_description = os.path.join(config["model_dir"], MODEL_DESCRIPTION_FILENAME)
 
     while max_batch_size - min_batch_size > min_range:
         batch_size = (max_batch_size + min_batch_size) // 2


Test patch:
diff --git a/opennmt/tests/config_test.py b/opennmt/tests/config_test.py
index 8ef4717c..586fb2eb 100644
--- a/opennmt/tests/config_test.py
+++ b/opennmt/tests/config_test.py
@@ -64,12 +64,14 @@ class ConfigTest(tf.test.TestCase):
         _check_model(
             config.load_model(model_dir, model_name=model_name, as_builder=as_builder)
         )
-        self.assertTrue(os.path.exists(os.path.join(model_dir, "model_description.py")))
+        self.assertTrue(
+            os.path.exists(os.path.join(model_dir, config.MODEL_DESCRIPTION_FILENAME))
+        )
         _check_model(config.load_model(model_dir, as_builder=as_builder))
 
     def testLoadModelDescriptionCompat(self):
         model_dir = self.get_temp_dir()
-        description = os.path.join(model_dir, "model_description.py")
+        description = os.path.join(model_dir, config.MODEL_DESCRIPTION_FILENAME)
         with open(description, "w") as description_file:
             description_file.write("from opennmt.models import catalog\n")
             description_file.write("model = catalog.Transformer\n")
@@ -80,7 +82,9 @@ class ConfigTest(tf.test.TestCase):
         model_file = self._writeCustomModel()
         model_dir = self.get_temp_dir()
         model = config.load_model(model_dir, model_file=model_file)
-        saved_description_path = os.path.join(model_dir, "model_description.py")
+        saved_description_path = os.path.join(
+            model_dir, config.MODEL_DESCRIPTION_FILENAME
+        )
         self.assertTrue(os.path.exists(saved_description_path))
         self.assertTrue(filecmp.cmp(model_file, saved_description_path))
         self.assertEqual(model, 42)
@@ -89,7 +93,9 @@ class ConfigTest(tf.test.TestCase):
 
     def testLoadModelFileOverride(self):
         model_dir = self.get_temp_dir()
-        saved_description_path = os.path.join(model_dir, "model_description.py")
+        saved_description_path = os.path.join(
+            model_dir, config.MODEL_DESCRIPTION_FILENAME
+        )
         model_file = self._writeCustomModel(filename="test_model1.py", return_value=1)
         config.load_model(model_dir, model_file=model_file)
         self.assertTrue(filecmp.cmp(model_file, saved_description_path))
diff --git a/opennmt/tests/runner_test.py b/opennmt/tests/runner_test.py
index 9f86d55f..7510e1f9 100644
--- a/opennmt/tests/runner_test.py
+++ b/opennmt/tests/runner_test.py
@@ -11,7 +11,7 @@ import tensorflow as tf
 from opennmt import decoders
 from opennmt import models
 from opennmt import Runner
-from opennmt.config import load_model
+from opennmt.config import load_model, MODEL_DESCRIPTION_FILENAME
 from opennmt.utils import exporters
 from opennmt.utils import misc
 from opennmt.tests import test_util
@@ -101,6 +101,9 @@ class RunnerTest(tf.test.TestCase):
         self.assertLen(
             tf.train.get_checkpoint_state(avg_dir).all_model_checkpoint_paths, 1
         )
+        self.assertTrue(
+            os.path.isfile(os.path.join(avg_dir, MODEL_DESCRIPTION_FILENAME))
+        )
         model_dir = os.path.dirname(avg_dir)
         self.assertEndsWith(tf.train.latest_checkpoint(model_dir), "145002")
         self.assertLen(
@@ -270,6 +273,9 @@ class RunnerTest(tf.test.TestCase):
             runner.update_vocab(output_dir, tgt_vocab=new_en_vocab), output_dir
         )
         self.assertEqual(runner.model_dir, output_dir)
+        self.assertTrue(
+            os.path.isfile(os.path.join(output_dir, MODEL_DESCRIPTION_FILENAME))
+        )
 
         # Check that the translation is unchanged.
         en_file = os.path.join(self.get_temp_dir(), "output.txt")