We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
DataLossError when infer with the model from checkpoints averaging
Hi,

tensorflow-gpu 1.8.0
opennmt-tf 1.5.0

Model is trained with the same parameter of [wmt_ende.yml](https://github.com/OpenNMT/OpenNMT-tf/blob/master/scripts/wmt/config/wmt_ende.yml) with 4 x 1080Ti

After the checkpoints averaging (max_count=5), 
```bash
onmt-average-checkpoints --max_count=5 --model_dir=model/ --output_dir=model/avg/
```
do the inference in averaged model, but it return an error,
```bash
onmt-main infer --model_type Transformer --config train_enfr.yml --features_file newstest2014-fren-src.en.sp --predictions_file newstest2014-fren-src.en.sp.pred --checkpoint_path /model/avg/
```

```bash
DataLossError (see above for traceback): Invalid size in bundle entry: key transformer/decoder/LayerNorm/beta; stored size 4096; expected size 2048
         [[Node: save/RestoreV2 = RestoreV2[dtypes=[DT_INT64, DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, ..., DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT], _device="/job:localhost/replica:0/task:0/device:CPU:0"](_arg_save/Const_0_0, save/RestoreV2/tensor_names, save/RestoreV2/shape_and_slices)]]
         [[Node: save/RestoreV2/_301 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device_incarnation=1, tensor_name="edge_306_save/RestoreV2", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:GPU:0"]()]]
``` 
But there is **no error** if I do inference with each single checkpoint, 
```bash
onmt-main infer --model_type Transformer --config train_enfr.yml --features_file newstest2014-fren-src.en.sp --predictions_file newstest2014-fren-src.en.sp.pred --checkpoint_path /model/model.ckpt-345921
```

I wonder if there something wrong when I average the checkpoints ?

Thanks !


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `opennmt/utils/checkpoint.py`
Content:
```
1 """Checkpoint utilities."""
2 
3 import os
4 import six
5 
6 import tensorflow as tf
7 import numpy as np
8 
9 from opennmt.utils.vocab import Vocab
10 from opennmt.utils.misc import count_lines
11 
12 
13 def _get_vocabulary_mapping(current_vocab_path, new_vocab_path, mode):
14   """Maps vocabulary new indices to old ones. -1 means that the entry is new."""
15   current_vocab = Vocab(from_file=current_vocab_path)
16   new_vocab = Vocab(from_file=new_vocab_path)
17   mapping = []
18   if mode == "merge":
19     mapping = [i for i in range(current_vocab.size)]
20     for new_word in new_vocab.words:
21       if current_vocab.lookup(new_word) is None:
22         mapping.append(-1)
23   elif mode == "replace":
24     for new_word in new_vocab.words:
25       idx = current_vocab.lookup(new_word)
26       if idx is not None:
27         mapping.append(idx)
28       else:
29         mapping.append(-1)
30   mapping.append(current_vocab.size)  # <unk> token is always the last entry.
31   return mapping
32 
33 def _update_vocabulary_variable(variable, vocab_size, mapping):
34   """Creates a new variable, possibly copying previous entries based on mapping."""
35   dim = variable.shape.index(vocab_size)
36   # Make the dimension to index the first.
37   perm = list(range(len(variable.shape)))
38   perm[0], perm[dim] = perm[dim], perm[0]
39   variable_t = np.transpose(variable, axes=perm)
40   new_shape = list(variable_t.shape)
41   new_shape[0] = len(mapping)
42   new_variable_t = np.zeros(new_shape, dtype=variable.dtype)
43   for i, j in enumerate(mapping):
44     if j >= 0:
45       new_variable_t[i] = variable_t[j]
46   new_variable = np.transpose(new_variable_t, axes=perm)
47   return new_variable
48 
49 def _update_vocabulary_variables(variables, current_vocab_path, new_vocab_path, scope_name, mode):
50   current_size = count_lines(current_vocab_path) + 1
51   mapping = _get_vocabulary_mapping(current_vocab_path, new_vocab_path, mode)
52   for name, tensor in six.iteritems(variables):
53     if scope_name in name and any(d == current_size for d in tensor.shape):
54       tf.logging.debug("Updating variable %s" % name)
55       variables[name] = _update_vocabulary_variable(tensor, current_size, mapping)
56 
57 def _save_new_variables(variables, output_dir, base_checkpoint_path, session_config=None):
58   if "global_step" in variables:
59     del variables["global_step"]
60   tf_vars = []
61   for name, value in six.iteritems(variables):
62     trainable = True
63     dtype = tf.as_dtype(value.dtype)
64     if name.startswith("words_per_sec"):
65       trainable = False
66       dtype = tf.int64  # TODO: why is the dtype not correct for these variables?
67     tf_vars.append(tf.get_variable(
68         name,
69         shape=value.shape,
70         dtype=dtype,
71         trainable=trainable))
72   placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
73   assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
74 
75   latest_step = int(base_checkpoint_path.split("-")[-1])
76   out_base_file = os.path.join(output_dir, "model.ckpt")
77   global_step = tf.get_variable(
78       "global_step",
79       initializer=tf.constant(latest_step, dtype=tf.int64),
80       trainable=False)
81   saver = tf.train.Saver(tf.global_variables())
82 
83   with tf.Session(config=session_config) as sess:
84     sess.run(tf.global_variables_initializer())
85     for p, assign_op, (name, value) in zip(placeholders, assign_ops, six.iteritems(variables)):
86       sess.run(assign_op, {p: value})
87     tf.logging.info("Saving new checkpoint to %s" % output_dir)
88     saver.save(sess, out_base_file, global_step=global_step)
89 
90   return output_dir
91 
92 def update_vocab(model_dir,
93                  output_dir,
94                  current_src_vocab,
95                  current_tgt_vocab,
96                  new_src_vocab=None,
97                  new_tgt_vocab=None,
98                  mode="merge",
99                  session_config=None):
100   """Updates the last checkpoint to support new vocabularies.
101 
102   This allows to add new words to a model while keeping the previously learned
103   weights.
104 
105   Args:
106     model_dir: The directory containing checkpoints (the most recent will be loaded).
107     output_dir: The directory that will contain the converted checkpoint.
108     current_src_vocab: Path to the source vocabulary currently use in the model.
109     current_tgt_vocab: Path to the target vocabulary currently use in the model.
110     new_src_vocab: Path to the new source vocabulary to support.
111     new_tgt_vocab: Path to the new target vocabulary to support.
112     mode: Update mode: "merge" keeps all existing words and adds new words,
113       "replace" makes the new vocabulary file the active one. In all modes,
114       if an existing word appears in the new vocabulary, its learned weights
115       are kept in the converted checkpoint.
116     session_config: Optional configuration to use when creating the session.
117 
118   Returns:
119     The path to the directory containing the converted checkpoint.
120 
121   Raises:
122     ValueError: if :obj:`output_dir` is the same as :obj:`model_dir` or if
123       :obj:`mode` is invalid.
124   """
125   if model_dir == output_dir:
126     raise ValueError("Model and output directory must be different")
127   if mode not in ("merge", "replace"):
128     raise ValueError("invalid vocab update mode: %s" % mode)
129   if new_src_vocab is None and new_tgt_vocab is None:
130     return model_dir
131   checkpoint_path = tf.train.latest_checkpoint(model_dir)
132   tf.logging.info("Updating vocabulary related variables in checkpoint %s" % checkpoint_path)
133   reader = tf.train.load_checkpoint(checkpoint_path)
134   variable_map = reader.get_variable_to_shape_map()
135   variable_value = {name:reader.get_tensor(name) for name, _ in six.iteritems(variable_map)}
136   if new_src_vocab is not None:
137     _update_vocabulary_variables(variable_value, current_src_vocab, new_src_vocab, "encoder", mode)
138   if new_tgt_vocab is not None:
139     _update_vocabulary_variables(variable_value, current_tgt_vocab, new_tgt_vocab, "decoder", mode)
140   return _save_new_variables(
141       variable_value,
142       output_dir,
143       checkpoint_path,
144       session_config=session_config)
145 
146 
147 def average_checkpoints(model_dir, output_dir, max_count=8, session_config=None):
148   """Averages checkpoints.
149 
150   Args:
151     model_dir: The directory containing checkpoints.
152     output_dir: The directory that will contain the averaged checkpoint.
153     max_count: The maximum number of checkpoints to average.
154     session_config: Configuration to use when creating the session.
155 
156   Returns:
157     The path to the directory containing the averaged checkpoint.
158 
159   Raises:
160     ValueError: if :obj:`output_dir` is the same as :obj:`model_dir`.
161   """
162   # This script is modified version of
163   # https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_avg_all.py
164   # which comes with the following license and copyright notice:
165 
166   # Copyright 2017 The Tensor2Tensor Authors.
167   #
168   # Licensed under the Apache License, Version 2.0 (the "License");
169   # you may not use this file except in compliance with the License.
170   # You may obtain a copy of the License at
171   #
172   #     http://www.apache.org/licenses/LICENSE-2.0
173   #
174   # Unless required by applicable law or agreed to in writing, software
175   # distributed under the License is distributed on an "AS IS" BASIS,
176   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
177   # See the License for the specific language governing permissions and
178   # limitations under the License.
179   if model_dir == output_dir:
180     raise ValueError("Model and output directory must be different")
181 
182   checkpoints_path = tf.train.get_checkpoint_state(model_dir).all_model_checkpoint_paths
183   if len(checkpoints_path) > max_count:
184     checkpoints_path = checkpoints_path[-max_count:]
185   num_checkpoints = len(checkpoints_path)
186 
187   tf.logging.info("Averaging %d checkpoints..." % num_checkpoints)
188   tf.logging.info("Listing variables...")
189 
190   var_list = tf.train.list_variables(checkpoints_path[0])
191   avg_values = {}
192   for name, shape in var_list:
193     if not name.startswith("global_step"):
194       avg_values[name] = np.zeros(shape)
195 
196   for checkpoint_path in checkpoints_path:
197     tf.logging.info("Loading checkpoint %s" % checkpoint_path)
198     reader = tf.train.load_checkpoint(checkpoint_path)
199     for name in avg_values:
200       avg_values[name] += reader.get_tensor(name) / num_checkpoints
201 
202   return _save_new_variables(
203       avg_values,
204       output_dir,
205       checkpoints_path[-1],
206       session_config=session_config)
207
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/opennmt/utils/checkpoint.py b/opennmt/utils/checkpoint.py
--- a/opennmt/utils/checkpoint.py
+++ b/opennmt/utils/checkpoint.py
@@ -191,7 +191,7 @@
   avg_values = {}
   for name, shape in var_list:
     if not name.startswith("global_step"):
-      avg_values[name] = np.zeros(shape)
+      avg_values[name] = np.zeros(shape, dtype=np.float32)
 
   for checkpoint_path in checkpoints_path:
     tf.logging.info("Loading checkpoint %s" % checkpoint_path)

{"golden_diff": "diff --git a/opennmt/utils/checkpoint.py b/opennmt/utils/checkpoint.py\n--- a/opennmt/utils/checkpoint.py\n+++ b/opennmt/utils/checkpoint.py\n@@ -191,7 +191,7 @@\n   avg_values = {}\n   for name, shape in var_list:\n     if not name.startswith(\"global_step\"):\n-      avg_values[name] = np.zeros(shape)\n+      avg_values[name] = np.zeros(shape, dtype=np.float32)\n \n   for checkpoint_path in checkpoints_path:\n     tf.logging.info(\"Loading checkpoint %s\" % checkpoint_path)\n", "issue": "DataLossError when infer with the model from checkpoints averaging\nHi,\r\n\r\ntensorflow-gpu 1.8.0\r\nopennmt-tf 1.5.0\r\n\r\nModel is trained with the same parameter of [wmt_ende.yml](https://github.com/OpenNMT/OpenNMT-tf/blob/master/scripts/wmt/config/wmt_ende.yml) with 4 x 1080Ti\r\n\r\nAfter the checkpoints averaging (max_count=5), \r\n```bash\r\nonmt-average-checkpoints --max_count=5 --model_dir=model/ --output_dir=model/avg/\r\n```\r\ndo the inference in averaged model, but it return an error,\r\n```bash\r\nonmt-main infer --model_type Transformer --config train_enfr.yml --features_file newstest2014-fren-src.en.sp --predictions_file newstest2014-fren-src.en.sp.pred --checkpoint_path /model/avg/\r\n```\r\n\r\n```bash\r\nDataLossError (see above for traceback): Invalid size in bundle entry: key transformer/decoder/LayerNorm/beta; stored size 4096; expected size 2048\r\n         [[Node: save/RestoreV2 = RestoreV2[dtypes=[DT_INT64, DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, ..., DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT], _device=\"/job:localhost/replica:0/task:0/device:CPU:0\"](_arg_save/Const_0_0, save/RestoreV2/tensor_names, save/RestoreV2/shape_and_slices)]]\r\n         [[Node: save/RestoreV2/_301 = _Recv[client_terminated=false, recv_device=\"/job:localhost/replica:0/task:0/device:GPU:0\", send_device=\"/job:localhost/replica:0/task:0/device:CPU:0\", send_device_incarnation=1, tensor_name=\"edge_306_save/RestoreV2\", tensor_type=DT_FLOAT, _device=\"/job:localhost/replica:0/task:0/device:GPU:0\"]()]]\r\n``` \r\nBut there is **no error** if I do inference with each single checkpoint, \r\n```bash\r\nonmt-main infer --model_type Transformer --config train_enfr.yml --features_file newstest2014-fren-src.en.sp --predictions_file newstest2014-fren-src.en.sp.pred --checkpoint_path /model/model.ckpt-345921\r\n```\r\n\r\nI wonder if there something wrong when I average the checkpoints ?\r\n\r\nThanks !\r\n\n", "before_files": [{"content": "\"\"\"Checkpoint utilities.\"\"\"\n\nimport os\nimport six\n\nimport tensorflow as tf\nimport numpy as np\n\nfrom opennmt.utils.vocab import Vocab\nfrom opennmt.utils.misc import count_lines\n\n\ndef _get_vocabulary_mapping(current_vocab_path, new_vocab_path, mode):\n  \"\"\"Maps vocabulary new indices to old ones. -1 means that the entry is new.\"\"\"\n  current_vocab = Vocab(from_file=current_vocab_path)\n  new_vocab = Vocab(from_file=new_vocab_path)\n  mapping = []\n  if mode == \"merge\":\n    mapping = [i for i in range(current_vocab.size)]\n    for new_word in new_vocab.words:\n      if current_vocab.lookup(new_word) is None:\n        mapping.append(-1)\n  elif mode == \"replace\":\n    for new_word in new_vocab.words:\n      idx = current_vocab.lookup(new_word)\n      if idx is not None:\n        mapping.append(idx)\n      else:\n        mapping.append(-1)\n  mapping.append(current_vocab.size)  # <unk> token is always the last entry.\n  return mapping\n\ndef _update_vocabulary_variable(variable, vocab_size, mapping):\n  \"\"\"Creates a new variable, possibly copying previous entries based on mapping.\"\"\"\n  dim = variable.shape.index(vocab_size)\n  # Make the dimension to index the first.\n  perm = list(range(len(variable.shape)))\n  perm[0], perm[dim] = perm[dim], perm[0]\n  variable_t = np.transpose(variable, axes=perm)\n  new_shape = list(variable_t.shape)\n  new_shape[0] = len(mapping)\n  new_variable_t = np.zeros(new_shape, dtype=variable.dtype)\n  for i, j in enumerate(mapping):\n    if j >= 0:\n      new_variable_t[i] = variable_t[j]\n  new_variable = np.transpose(new_variable_t, axes=perm)\n  return new_variable\n\ndef _update_vocabulary_variables(variables, current_vocab_path, new_vocab_path, scope_name, mode):\n  current_size = count_lines(current_vocab_path) + 1\n  mapping = _get_vocabulary_mapping(current_vocab_path, new_vocab_path, mode)\n  for name, tensor in six.iteritems(variables):\n    if scope_name in name and any(d == current_size for d in tensor.shape):\n      tf.logging.debug(\"Updating variable %s\" % name)\n      variables[name] = _update_vocabulary_variable(tensor, current_size, mapping)\n\ndef _save_new_variables(variables, output_dir, base_checkpoint_path, session_config=None):\n  if \"global_step\" in variables:\n    del variables[\"global_step\"]\n  tf_vars = []\n  for name, value in six.iteritems(variables):\n    trainable = True\n    dtype = tf.as_dtype(value.dtype)\n    if name.startswith(\"words_per_sec\"):\n      trainable = False\n      dtype = tf.int64  # TODO: why is the dtype not correct for these variables?\n    tf_vars.append(tf.get_variable(\n        name,\n        shape=value.shape,\n        dtype=dtype,\n        trainable=trainable))\n  placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]\n  assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]\n\n  latest_step = int(base_checkpoint_path.split(\"-\")[-1])\n  out_base_file = os.path.join(output_dir, \"model.ckpt\")\n  global_step = tf.get_variable(\n      \"global_step\",\n      initializer=tf.constant(latest_step, dtype=tf.int64),\n      trainable=False)\n  saver = tf.train.Saver(tf.global_variables())\n\n  with tf.Session(config=session_config) as sess:\n    sess.run(tf.global_variables_initializer())\n    for p, assign_op, (name, value) in zip(placeholders, assign_ops, six.iteritems(variables)):\n      sess.run(assign_op, {p: value})\n    tf.logging.info(\"Saving new checkpoint to %s\" % output_dir)\n    saver.save(sess, out_base_file, global_step=global_step)\n\n  return output_dir\n\ndef update_vocab(model_dir,\n                 output_dir,\n                 current_src_vocab,\n                 current_tgt_vocab,\n                 new_src_vocab=None,\n                 new_tgt_vocab=None,\n                 mode=\"merge\",\n                 session_config=None):\n  \"\"\"Updates the last checkpoint to support new vocabularies.\n\n  This allows to add new words to a model while keeping the previously learned\n  weights.\n\n  Args:\n    model_dir: The directory containing checkpoints (the most recent will be loaded).\n    output_dir: The directory that will contain the converted checkpoint.\n    current_src_vocab: Path to the source vocabulary currently use in the model.\n    current_tgt_vocab: Path to the target vocabulary currently use in the model.\n    new_src_vocab: Path to the new source vocabulary to support.\n    new_tgt_vocab: Path to the new target vocabulary to support.\n    mode: Update mode: \"merge\" keeps all existing words and adds new words,\n      \"replace\" makes the new vocabulary file the active one. In all modes,\n      if an existing word appears in the new vocabulary, its learned weights\n      are kept in the converted checkpoint.\n    session_config: Optional configuration to use when creating the session.\n\n  Returns:\n    The path to the directory containing the converted checkpoint.\n\n  Raises:\n    ValueError: if :obj:`output_dir` is the same as :obj:`model_dir` or if\n      :obj:`mode` is invalid.\n  \"\"\"\n  if model_dir == output_dir:\n    raise ValueError(\"Model and output directory must be different\")\n  if mode not in (\"merge\", \"replace\"):\n    raise ValueError(\"invalid vocab update mode: %s\" % mode)\n  if new_src_vocab is None and new_tgt_vocab is None:\n    return model_dir\n  checkpoint_path = tf.train.latest_checkpoint(model_dir)\n  tf.logging.info(\"Updating vocabulary related variables in checkpoint %s\" % checkpoint_path)\n  reader = tf.train.load_checkpoint(checkpoint_path)\n  variable_map = reader.get_variable_to_shape_map()\n  variable_value = {name:reader.get_tensor(name) for name, _ in six.iteritems(variable_map)}\n  if new_src_vocab is not None:\n    _update_vocabulary_variables(variable_value, current_src_vocab, new_src_vocab, \"encoder\", mode)\n  if new_tgt_vocab is not None:\n    _update_vocabulary_variables(variable_value, current_tgt_vocab, new_tgt_vocab, \"decoder\", mode)\n  return _save_new_variables(\n      variable_value,\n      output_dir,\n      checkpoint_path,\n      session_config=session_config)\n\n\ndef average_checkpoints(model_dir, output_dir, max_count=8, session_config=None):\n  \"\"\"Averages checkpoints.\n\n  Args:\n    model_dir: The directory containing checkpoints.\n    output_dir: The directory that will contain the averaged checkpoint.\n    max_count: The maximum number of checkpoints to average.\n    session_config: Configuration to use when creating the session.\n\n  Returns:\n    The path to the directory containing the averaged checkpoint.\n\n  Raises:\n    ValueError: if :obj:`output_dir` is the same as :obj:`model_dir`.\n  \"\"\"\n  # This script is modified version of\n  # https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_avg_all.py\n  # which comes with the following license and copyright notice:\n\n  # Copyright 2017 The Tensor2Tensor Authors.\n  #\n  # Licensed under the Apache License, Version 2.0 (the \"License\");\n  # you may not use this file except in compliance with the License.\n  # You may obtain a copy of the License at\n  #\n  #     http://www.apache.org/licenses/LICENSE-2.0\n  #\n  # Unless required by applicable law or agreed to in writing, software\n  # distributed under the License is distributed on an \"AS IS\" BASIS,\n  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n  # See the License for the specific language governing permissions and\n  # limitations under the License.\n  if model_dir == output_dir:\n    raise ValueError(\"Model and output directory must be different\")\n\n  checkpoints_path = tf.train.get_checkpoint_state(model_dir).all_model_checkpoint_paths\n  if len(checkpoints_path) > max_count:\n    checkpoints_path = checkpoints_path[-max_count:]\n  num_checkpoints = len(checkpoints_path)\n\n  tf.logging.info(\"Averaging %d checkpoints...\" % num_checkpoints)\n  tf.logging.info(\"Listing variables...\")\n\n  var_list = tf.train.list_variables(checkpoints_path[0])\n  avg_values = {}\n  for name, shape in var_list:\n    if not name.startswith(\"global_step\"):\n      avg_values[name] = np.zeros(shape)\n\n  for checkpoint_path in checkpoints_path:\n    tf.logging.info(\"Loading checkpoint %s\" % checkpoint_path)\n    reader = tf.train.load_checkpoint(checkpoint_path)\n    for name in avg_values:\n      avg_values[name] += reader.get_tensor(name) / num_checkpoints\n\n  return _save_new_variables(\n      avg_values,\n      output_dir,\n      checkpoints_path[-1],\n      session_config=session_config)\n", "path": "opennmt/utils/checkpoint.py"}], "after_files": [{"content": "\"\"\"Checkpoint utilities.\"\"\"\n\nimport os\nimport six\n\nimport tensorflow as tf\nimport numpy as np\n\nfrom opennmt.utils.vocab import Vocab\nfrom opennmt.utils.misc import count_lines\n\n\ndef _get_vocabulary_mapping(current_vocab_path, new_vocab_path, mode):\n  \"\"\"Maps vocabulary new indices to old ones. -1 means that the entry is new.\"\"\"\n  current_vocab = Vocab(from_file=current_vocab_path)\n  new_vocab = Vocab(from_file=new_vocab_path)\n  mapping = []\n  if mode == \"merge\":\n    mapping = [i for i in range(current_vocab.size)]\n    for new_word in new_vocab.words:\n      if current_vocab.lookup(new_word) is None:\n        mapping.append(-1)\n  elif mode == \"replace\":\n    for new_word in new_vocab.words:\n      idx = current_vocab.lookup(new_word)\n      if idx is not None:\n        mapping.append(idx)\n      else:\n        mapping.append(-1)\n  mapping.append(current_vocab.size)  # <unk> token is always the last entry.\n  return mapping\n\ndef _update_vocabulary_variable(variable, vocab_size, mapping):\n  \"\"\"Creates a new variable, possibly copying previous entries based on mapping.\"\"\"\n  dim = variable.shape.index(vocab_size)\n  # Make the dimension to index the first.\n  perm = list(range(len(variable.shape)))\n  perm[0], perm[dim] = perm[dim], perm[0]\n  variable_t = np.transpose(variable, axes=perm)\n  new_shape = list(variable_t.shape)\n  new_shape[0] = len(mapping)\n  new_variable_t = np.zeros(new_shape, dtype=variable.dtype)\n  for i, j in enumerate(mapping):\n    if j >= 0:\n      new_variable_t[i] = variable_t[j]\n  new_variable = np.transpose(new_variable_t, axes=perm)\n  return new_variable\n\ndef _update_vocabulary_variables(variables, current_vocab_path, new_vocab_path, scope_name, mode):\n  current_size = count_lines(current_vocab_path) + 1\n  mapping = _get_vocabulary_mapping(current_vocab_path, new_vocab_path, mode)\n  for name, tensor in six.iteritems(variables):\n    if scope_name in name and any(d == current_size for d in tensor.shape):\n      tf.logging.debug(\"Updating variable %s\" % name)\n      variables[name] = _update_vocabulary_variable(tensor, current_size, mapping)\n\ndef _save_new_variables(variables, output_dir, base_checkpoint_path, session_config=None):\n  if \"global_step\" in variables:\n    del variables[\"global_step\"]\n  tf_vars = []\n  for name, value in six.iteritems(variables):\n    trainable = True\n    dtype = tf.as_dtype(value.dtype)\n    if name.startswith(\"words_per_sec\"):\n      trainable = False\n      dtype = tf.int64  # TODO: why is the dtype not correct for these variables?\n    tf_vars.append(tf.get_variable(\n        name,\n        shape=value.shape,\n        dtype=dtype,\n        trainable=trainable))\n  placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]\n  assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]\n\n  latest_step = int(base_checkpoint_path.split(\"-\")[-1])\n  out_base_file = os.path.join(output_dir, \"model.ckpt\")\n  global_step = tf.get_variable(\n      \"global_step\",\n      initializer=tf.constant(latest_step, dtype=tf.int64),\n      trainable=False)\n  saver = tf.train.Saver(tf.global_variables())\n\n  with tf.Session(config=session_config) as sess:\n    sess.run(tf.global_variables_initializer())\n    for p, assign_op, (name, value) in zip(placeholders, assign_ops, six.iteritems(variables)):\n      sess.run(assign_op, {p: value})\n    tf.logging.info(\"Saving new checkpoint to %s\" % output_dir)\n    saver.save(sess, out_base_file, global_step=global_step)\n\n  return output_dir\n\ndef update_vocab(model_dir,\n                 output_dir,\n                 current_src_vocab,\n                 current_tgt_vocab,\n                 new_src_vocab=None,\n                 new_tgt_vocab=None,\n                 mode=\"merge\",\n                 session_config=None):\n  \"\"\"Updates the last checkpoint to support new vocabularies.\n\n  This allows to add new words to a model while keeping the previously learned\n  weights.\n\n  Args:\n    model_dir: The directory containing checkpoints (the most recent will be loaded).\n    output_dir: The directory that will contain the converted checkpoint.\n    current_src_vocab: Path to the source vocabulary currently use in the model.\n    current_tgt_vocab: Path to the target vocabulary currently use in the model.\n    new_src_vocab: Path to the new source vocabulary to support.\n    new_tgt_vocab: Path to the new target vocabulary to support.\n    mode: Update mode: \"merge\" keeps all existing words and adds new words,\n      \"replace\" makes the new vocabulary file the active one. In all modes,\n      if an existing word appears in the new vocabulary, its learned weights\n      are kept in the converted checkpoint.\n    session_config: Optional configuration to use when creating the session.\n\n  Returns:\n    The path to the directory containing the converted checkpoint.\n\n  Raises:\n    ValueError: if :obj:`output_dir` is the same as :obj:`model_dir` or if\n      :obj:`mode` is invalid.\n  \"\"\"\n  if model_dir == output_dir:\n    raise ValueError(\"Model and output directory must be different\")\n  if mode not in (\"merge\", \"replace\"):\n    raise ValueError(\"invalid vocab update mode: %s\" % mode)\n  if new_src_vocab is None and new_tgt_vocab is None:\n    return model_dir\n  checkpoint_path = tf.train.latest_checkpoint(model_dir)\n  tf.logging.info(\"Updating vocabulary related variables in checkpoint %s\" % checkpoint_path)\n  reader = tf.train.load_checkpoint(checkpoint_path)\n  variable_map = reader.get_variable_to_shape_map()\n  variable_value = {name:reader.get_tensor(name) for name, _ in six.iteritems(variable_map)}\n  if new_src_vocab is not None:\n    _update_vocabulary_variables(variable_value, current_src_vocab, new_src_vocab, \"encoder\", mode)\n  if new_tgt_vocab is not None:\n    _update_vocabulary_variables(variable_value, current_tgt_vocab, new_tgt_vocab, \"decoder\", mode)\n  return _save_new_variables(\n      variable_value,\n      output_dir,\n      checkpoint_path,\n      session_config=session_config)\n\n\ndef average_checkpoints(model_dir, output_dir, max_count=8, session_config=None):\n  \"\"\"Averages checkpoints.\n\n  Args:\n    model_dir: The directory containing checkpoints.\n    output_dir: The directory that will contain the averaged checkpoint.\n    max_count: The maximum number of checkpoints to average.\n    session_config: Configuration to use when creating the session.\n\n  Returns:\n    The path to the directory containing the averaged checkpoint.\n\n  Raises:\n    ValueError: if :obj:`output_dir` is the same as :obj:`model_dir`.\n  \"\"\"\n  # This script is modified version of\n  # https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/bin/t2t_avg_all.py\n  # which comes with the following license and copyright notice:\n\n  # Copyright 2017 The Tensor2Tensor Authors.\n  #\n  # Licensed under the Apache License, Version 2.0 (the \"License\");\n  # you may not use this file except in compliance with the License.\n  # You may obtain a copy of the License at\n  #\n  #     http://www.apache.org/licenses/LICENSE-2.0\n  #\n  # Unless required by applicable law or agreed to in writing, software\n  # distributed under the License is distributed on an \"AS IS\" BASIS,\n  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n  # See the License for the specific language governing permissions and\n  # limitations under the License.\n  if model_dir == output_dir:\n    raise ValueError(\"Model and output directory must be different\")\n\n  checkpoints_path = tf.train.get_checkpoint_state(model_dir).all_model_checkpoint_paths\n  if len(checkpoints_path) > max_count:\n    checkpoints_path = checkpoints_path[-max_count:]\n  num_checkpoints = len(checkpoints_path)\n\n  tf.logging.info(\"Averaging %d checkpoints...\" % num_checkpoints)\n  tf.logging.info(\"Listing variables...\")\n\n  var_list = tf.train.list_variables(checkpoints_path[0])\n  avg_values = {}\n  for name, shape in var_list:\n    if not name.startswith(\"global_step\"):\n      avg_values[name] = np.zeros(shape, dtype=np.float32)\n\n  for checkpoint_path in checkpoints_path:\n    tf.logging.info(\"Loading checkpoint %s\" % checkpoint_path)\n    reader = tf.train.load_checkpoint(checkpoint_path)\n    for name in avg_values:\n      avg_values[name] += reader.get_tensor(name) / num_checkpoints\n\n  return _save_new_variables(\n      avg_values,\n      output_dir,\n      checkpoints_path[-1],\n      session_config=session_config)\n", "path": "opennmt/utils/checkpoint.py"}]}