Repository: OpenNMT/OpenNMT-tf

Issue:
Checkpoint conversion between float16 and float32
Users should be able to convert a model trained with `float16` to `float32` (or vice versa).

* Add an API endpoint in `opennmt.utils.checkpoint`
* Add a CLI tool in `opennmt.bin.convert_checkpoint`


Discussion:

Resolving patch:
diff --git a/opennmt/bin/convert_checkpoint.py b/opennmt/bin/convert_checkpoint.py
new file mode 100644
--- /dev/null
+++ b/opennmt/bin/convert_checkpoint.py
@@ -0,0 +1,44 @@
+"""Script to convert checkpoint variables from one data type to another."""
+
+import argparse
+
+import tensorflow as tf
+
+from opennmt.utils import checkpoint
+
+
+def main():
+  tf.logging.set_verbosity(tf.logging.INFO)
+
+  parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument("--model_dir", default=None,
+                      help="The path to the model directory.")
+  parser.add_argument("--checkpoint_path", default=None,
+                      help="The path to the checkpoint to convert.")
+  parser.add_argument("--output_dir", required=True,
+                      help="The output directory where the updated checkpoint will be saved.")
+  parser.add_argument("--target_dtype", required=True,
+                      help="Target data type (e.g. float16 or float32).")
+  parser.add_argument("--source_dtype", default=None,
+                      help="Source data type (e.g. float16 or float32, inferred if not set).")
+  args = parser.parse_args()
+  if args.model_dir is None and args.checkpoint_path is None:
+    raise ValueError("One of --checkpoint_path and --model_dir should be set")
+  checkpoint_path = args.checkpoint_path
+  if checkpoint_path is None:
+    checkpoint_path = tf.train.latest_checkpoint(args.model_dir)
+  target_dtype = tf.as_dtype(args.target_dtype)
+  if args.source_dtype is None:
+    source_dtype = tf.float32 if target_dtype == tf.float16 else tf.float16
+  else:
+    source_dtype = tf.as_dtype(args.source_dtype)
+  checkpoint.convert_checkpoint(
+      checkpoint_path,
+      args.output_dir,
+      source_dtype,
+      target_dtype,
+      session_config=tf.ConfigProto(device_count={"GPU": 0}))
+
+
+if __name__ == "__main__":
+  main()
diff --git a/opennmt/utils/checkpoint.py b/opennmt/utils/checkpoint.py
--- a/opennmt/utils/checkpoint.py
+++ b/opennmt/utils/checkpoint.py
@@ -54,8 +54,9 @@ def _update_vocabulary_variables(variables, current_vocab_path, new_vocab_path,
       tf.logging.debug("Updating variable %s" % name)
       variables[name] = _update_vocabulary_variable(tensor, current_size, mapping)
 
-def _save_new_variables(variables, output_dir, base_checkpoint_path, session_config=None):
+def _create_checkpoint_from_variables(variables, output_dir, latest_step=None, session_config=None):
   if "global_step" in variables:
+    latest_step = variables["global_step"]
     del variables["global_step"]
   tf_vars = []
   for name, value in six.iteritems(variables):
@@ -72,7 +73,6 @@ def _save_new_variables(variables, output_dir, base_checkpoint_path, session_con
   placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
   assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
 
-  latest_step = int(base_checkpoint_path.split("-")[-1])
   out_base_file = os.path.join(output_dir, "model.ckpt")
   global_step = tf.get_variable(
       "global_step",
@@ -89,6 +89,52 @@ def _save_new_variables(variables, output_dir, base_checkpoint_path, session_con
 
   return output_dir
 
+def get_checkpoint_variables(checkpoint_path):
+  """Returns variables included in a checkpoint.
+
+  Args:
+    checkpoint_path: Path to the checkpoint.
+
+  Returns:
+    A dictionary mapping variables name to value.
+  """
+  reader = tf.train.load_checkpoint(checkpoint_path)
+  return {
+      name:reader.get_tensor(name)
+      for name in six.iterkeys(reader.get_variable_to_shape_map())}
+
+def convert_checkpoint(checkpoint_path,
+                       output_dir,
+                       source_dtype,
+                       target_type,
+                       session_config=None):
+  """Converts checkpoint variables from one dtype to another.
+
+  Args:
+    checkpoint_path: The path to the checkpoint to convert.
+    output_dir: The directory that will contain the converted checkpoint.
+    source_dtype: The data type to convert from.
+    target_dtype: The data type to convert to.
+    session_config: Optional configuration to use when creating the session.
+
+  Returns:
+    The path to the directory containing the converted checkpoint.
+
+  Raises:
+    ValueError: if :obj:`output_dir` points to the same directory as
+      :obj:`checkpoint_path`.
+  """
+  if os.path.dirname(checkpoint_path) == output_dir:
+    raise ValueError("Checkpoint and output directory must be different")
+  variables = get_checkpoint_variables(checkpoint_path)
+  for name, value in six.iteritems(variables):
+    if not name.startswith("optim") and tf.as_dtype(value.dtype) == source_dtype:
+      variables[name] = value.astype(target_type.as_numpy_dtype())
+  return _create_checkpoint_from_variables(
+      variables,
+      output_dir,
+      session_config=session_config)
+
 def update_vocab(model_dir,
                  output_dir,
                  current_src_vocab,
@@ -130,17 +176,14 @@ def update_vocab(model_dir,
     return model_dir
   checkpoint_path = tf.train.latest_checkpoint(model_dir)
   tf.logging.info("Updating vocabulary related variables in checkpoint %s" % checkpoint_path)
-  reader = tf.train.load_checkpoint(checkpoint_path)
-  variable_map = reader.get_variable_to_shape_map()
-  variable_value = {name:reader.get_tensor(name) for name, _ in six.iteritems(variable_map)}
+  variable_value = get_checkpoint_variables(checkpoint_path)
   if new_src_vocab is not None:
     _update_vocabulary_variables(variable_value, current_src_vocab, new_src_vocab, "encoder", mode)
   if new_tgt_vocab is not None:
     _update_vocabulary_variables(variable_value, current_tgt_vocab, new_tgt_vocab, "decoder", mode)
-  return _save_new_variables(
+  return _create_checkpoint_from_variables(
       variable_value,
       output_dir,
-      checkpoint_path,
       session_config=session_config)
 
 
@@ -199,8 +242,9 @@ def average_checkpoints(model_dir, output_dir, max_count=8, session_config=None)
     for name in avg_values:
       avg_values[name] += reader.get_tensor(name) / num_checkpoints
 
-  return _save_new_variables(
+  latest_step = int(checkpoints_path[-1].split("-")[-1])
+  return _create_checkpoint_from_variables(
       avg_values,
       output_dir,
-      checkpoints_path[-1],
+      latest_step=latest_step,
       session_config=session_config)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -55,6 +55,7 @@
             "onmt-ark-to-records=opennmt.bin.ark_to_records:main",
             "onmt-average-checkpoints=opennmt.bin.average_checkpoints:main",
             "onmt-build-vocab=opennmt.bin.build_vocab:main",
+            "onmt-convert-checkpoint=opennmt.bin.convert_checkpoint:main",
             "onmt-detokenize-text=opennmt.bin.detokenize_text:main",
             "onmt-main=opennmt.bin.main:main",
             "onmt-merge-config=opennmt.bin.merge_config:main",


Test patch:
diff --git a/opennmt/tests/checkpoint_test.py b/opennmt/tests/checkpoint_test.py
--- a/opennmt/tests/checkpoint_test.py
+++ b/opennmt/tests/checkpoint_test.py
@@ -49,14 +49,17 @@ def _generateCheckpoint(self,
     with tf.Graph().as_default() as graph:
       for name, value in six.iteritems(variables):
         if isinstance(value, tuple):
+          dtype = None
           initializer = tf.random_uniform_initializer()
           shape = value
         else:
-          initializer = tf.constant_initializer(value, dtype=tf.as_dtype(value.dtype))
+          dtype = tf.as_dtype(value.dtype)
+          initializer = tf.constant_initializer(value, dtype=dtype)
           shape = value.shape
         _ = tf.get_variable(
             name,
             shape=shape,
+            dtype=dtype,
             initializer=initializer)
       global_step = tf.get_variable(
           "global_step",
@@ -70,14 +73,6 @@ def _generateCheckpoint(self,
         saver.save(sess, os.path.join(model_dir, prefix), global_step=global_step)
       return saver.last_checkpoints[0], time.time()
 
-  def _readCheckpoint(model_dir, checkpoint_path=None):
-    if checkpoint_path is None:
-      checkpoint_path = tf.train.latest_checkpoint(model_dir)
-    reader = tf.train.load_checkpoint(checkpoint_path)
-    variable_map = reader.get_variable_to_shape_map()
-    variables = {name:reader.get_tensor(name) for name, _ in six.iteritems(variable_map)}
-    return variables
-
   def testCheckpointAveraging(self):
     model_dir = os.path.join(self.get_temp_dir(), "ckpt")
     os.makedirs(model_dir)
@@ -88,10 +83,27 @@ def testCheckpointAveraging(self):
         model_dir, 20, {"x": np.ones((2, 3), dtype=np.float32)}, last_checkpoints=checkpoints))
     avg_dir = os.path.join(model_dir, "avg")
     checkpoint.average_checkpoints(model_dir, avg_dir)
-    avg_var = self._readCheckpoint(avg_dir)
+    avg_var = checkpoint.get_checkpoint_variables(avg_dir)
     self.assertEqual(avg_var["global_step"], 20)
     self.assertAllEqual(avg_var["x"], np.full((2, 3), 0.5, dtype=np.float32))
 
+  def testCheckpointDTypeConversion(self):
+    model_dir = os.path.join(self.get_temp_dir(), "ckpt-fp32")
+    os.makedirs(model_dir)
+    variables = {
+      "x": np.ones((2, 3), dtype=np.float32),
+      "optim/x": np.ones((2, 3), dtype=np.float32),
+      "counter": np.int64(42)
+    }
+    checkpoint_path, _ = self._generateCheckpoint(model_dir, 10, variables)
+    half_dir = os.path.join(model_dir, "fp16")
+    checkpoint.convert_checkpoint(checkpoint_path, half_dir, tf.float32, tf.float16)
+    half_var = checkpoint.get_checkpoint_variables(half_dir)
+    self.assertEqual(half_var["global_step"], 10)
+    self.assertEqual(half_var["x"].dtype, np.float16)
+    self.assertEqual(half_var["optim/x"].dtype, np.float32)
+    self.assertEqual(half_var["counter"].dtype, np.int64)
+
 
 if __name__ == "__main__":
   tf.test.main()