We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Fine-tuning only a subset of parameters
Hi,

Is there the possibility of optimizing only a subset of parameters in the computation graph?

--- 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/optim.py`
Content:
```
1 """Optimization related functions."""
2 
3 import collections
4 
5 import tensorflow as tf
6 
7 from opennmt import optimizers
8 from opennmt.utils import decay
9 from opennmt.utils import misc
10 
11 
12 def learning_rate_decay_fn(decay_type,
13                            decay_rate,
14                            decay_steps,
15                            decay_step_duration=1,
16                            staircase=True,
17                            start_decay_steps=0,
18                            minimum_learning_rate=0):
19   """Returns the learning rate decay functions.
20 
21   Args:
22     decay_type: The type of decay. A function from ``tf.train`` or
23      :mod:`opennmt.utils.decay` as a string.
24     decay_rate: The decay rate to apply.
25     decay_steps: The decay steps as described in the decay type function.
26     decay_step_duration: The number of training steps that make 1 decay step.
27     staircase: If ``True``, learning rate is decayed in a staircase fashion.
28     start_decay_steps: Start decay after this many steps.
29     minimum_learning_rate: Do not decay past this learning rate value.
30 
31   Returns:
32     A function with signature
33     ``(learning_rate, global_step) -> decayed_learning_rate``.
34 
35   Raises:
36     ValueError: if :obj:`decay_type` can not be resolved.
37   """
38   decay_params = {
39       "decay_rate": decay_rate,
40       "decay_steps": decay_steps,
41       "staircase": staircase
42   }
43   return learning_rate_decay_fn_v2(
44       decay_type,
45       decay_params=decay_params,
46       decay_step_duration=decay_step_duration,
47       start_decay_step=start_decay_steps,
48       minimum_learning_rate=minimum_learning_rate)
49 
50 def learning_rate_decay_fn_v2(decay_type,
51                               decay_params=None,
52                               decay_step_duration=1,
53                               start_decay_step=0,
54                               minimum_learning_rate=0.0):
55   """Returns the learning rate decay function.
56 
57   Args:
58     decay_type: The type of decay. A function from ``tf.train`` or
59       :mod:`opennmt.utils.decay` as a string.
60     decay_params: Additional parameters for the decay function.
61     decay_step_duration: The number of training steps that make 1 decay step.
62     start_decay_step: Start decay after this many steps.
63     minimum_learning_rate: Do not decay past this learning rate value.
64 
65   Returns:
66     A function with signature
67     ``(learning_rate, global_step) -> decayed_learning_rate``.
68 
69   Raises:
70     ValueError: if :obj:`decay_type` can not be resolved.
71   """
72   if decay_params is None:
73     decay_params = {}
74 
75   def _decay_fn(learning_rate, global_step):
76     decay_op_name = None
77 
78     if decay_op_name is None:
79       decay_op_name = getattr(tf.train, decay_type, None)
80     if decay_op_name is None:
81       decay_op_name = getattr(decay, decay_type, None)
82     if decay_op_name is None:
83       raise ValueError("Unknown decay function: {}".format(decay_type))
84 
85     # Map the training step to a decay step.
86     step = tf.maximum(global_step - start_decay_step, 0)
87     step //= decay_step_duration
88 
89     learning_rate = decay_op_name(learning_rate, step, **decay_params)
90     return tf.maximum(learning_rate, minimum_learning_rate)
91 
92   return _decay_fn
93 
94 def get_optimizer_class(classname):
95   """Returns the optimizer class.
96 
97   Args:
98     classname: The name of the optimizer class in ``tf.train``,
99       ``tf.contrib.opt``, or ``opennmt.optimizers`` as a string.
100 
101   Returns:
102     A class inheriting from ``tf.train.Optimizer``.
103 
104   Raises:
105     ValueError: if :obj:`classname` can not be resolved.
106   """
107   optimizer_class = None
108 
109   if optimizer_class is None:
110     optimizer_class = getattr(tf.train, classname, None)
111   if optimizer_class is None:
112     optimizer_class = getattr(tf.contrib.opt, classname, None)
113   if optimizer_class is None:
114     optimizer_class = getattr(optimizers, classname, None)
115   if optimizer_class is None:
116     raise ValueError("Unknown optimizer class: {}".format(classname))
117 
118   return optimizer_class
119 
120 def optimize(*args, **kwargs):
121   """Wrapper around ``optimize_loss`` for backward compatibility."""
122   update_op, _ = optimize_loss(*args, **kwargs)
123   return update_op
124 
125 def optimize_loss(loss, params, mixed_precision=False, var_list=None, hvd=None):
126   """Minimizes the loss.
127 
128   Args:
129     loss: The loss to minimize.
130     params: A dictionary of hyperparameters.
131     mixed_precision: If ``True``, wraps the optimizer to maintain a float32 copy
132       of the weights.
133     var_list: The variables to update.
134     hvd: Optional Horovod object.
135 
136   Returns:
137     The loss minimization op and a list of internal variables to initialize.
138   """
139   regularization = params.get("regularization")
140   if regularization is not None:
141     loss += regularization_penalty(
142         regularization["type"], regularization["scale"], weights_list=var_list)
143 
144   global_step = tf.train.get_or_create_global_step()
145   with tf.variable_scope("optim"):
146     learning_rate = tf.constant(params["learning_rate"], dtype=tf.float32)
147     if params.get("decay_type") is not None:
148       decay_params = params.get("decay_params", {})
149       if "decay_rate" in params:
150         # Backward compatibility: fill params from previous options.
151         decay_params["decay_rate"] = params["decay_rate"]
152         decay_params["decay_steps"] = params["decay_steps"]
153         decay_params["staircase"] = params.get("staircase", True)
154       decay_fn = learning_rate_decay_fn_v2(
155           params["decay_type"],
156           decay_params=decay_params,
157           decay_step_duration=params.get("decay_step_duration", 1),
158           start_decay_step=params.get("start_decay_steps", 0),
159           minimum_learning_rate=params.get("minimum_learning_rate", 0))
160       learning_rate = decay_fn(learning_rate, global_step)
161     tf.summary.scalar("learning_rate", learning_rate)
162 
163     # Optimizer.
164     optimizer_class = get_optimizer_class(params["optimizer"])
165     optimizer_params = params.get("optimizer_params", {})
166     if optimizer_class.__name__ == "AdafactorOptimizer":
167       optimizer = optimizers.get_adafactor_optimizer_from_params(
168           optimizer_class, optimizer_params, learning_rate=learning_rate)
169     else:
170       weight_decay = params.get("weight_decay")
171       if weight_decay is not None:
172         optimizer_class = tf.contrib.opt.extend_with_decoupled_weight_decay(optimizer_class)
173         optimizer = optimizer_class(
174             weight_decay, learning_rate=learning_rate, **optimizer_params)
175       else:
176         optimizer = optimizer_class(learning_rate, **optimizer_params)
177     if mixed_precision:
178       from opennmt.optimizers.mixed_precision_wrapper import get_loss_scale_from_params
179       optimizer = optimizers.MixedPrecisionOptimizerWrapper(
180           optimizer, loss_scale=get_loss_scale_from_params(params))
181     if hvd is not None:
182       from opennmt.optimizers.distributed_optimizer import DistributedOptimizer
183       optimizer = DistributedOptimizer.from_params(optimizer, params=params.get("horovod"))
184 
185     # Gradients.
186     gradients = optimizer.compute_gradients(
187         loss, var_list=var_list, colocate_gradients_with_ops=True)
188     _summarize_gradients_norm("global_norm/gradient_norm", gradients)
189     if params.get("clip_gradients") is not None:
190       gradients = _clip_gradients_by_norm(gradients, float(params["clip_gradients"]))
191       _summarize_gradients_norm("global_norm/clipped_gradient_norm", gradients)
192 
193     return delayed_update(
194         optimizer,
195         gradients,
196         global_step,
197         accum_count=params.get("gradients_accum", 1))
198 
199 def delayed_update(optimizer, grads_and_vars, global_step, accum_count=1):
200   """Possibly delays the parameters update by first accumulating gradients.
201 
202   Args:
203     optimizer: The optimizer.
204     grads_and_vars: List of (gradient, variable) pairs.
205     global_step: The training step that will be increased when the parameters
206       are updated.
207     accum_count: The number of times to accumulate gradients, as a constant or
208       a ``tf.Tensor``.
209 
210   Returns:
211     An operation that conditionally applies the gradients and a list of internal
212     variables to initialize.
213   """
214 
215   def _apply_gradients(grads_and_vars, global_step=None):
216     if hasattr(optimizer, "allreduce_gradients"):
217       grads_and_vars = optimizer.allreduce_gradients(grads_and_vars)
218     kwargs = {}
219     if "decay_var_list" in misc.function_args(optimizer.apply_gradients):
220       kwargs["decay_var_list"] = [var for _, var in grads_and_vars if not _is_bias(var)]
221     return optimizer.apply_gradients(grads_and_vars, global_step=global_step, **kwargs)
222 
223   if not tf.contrib.framework.is_tensor(accum_count) and accum_count == 1:
224     return _apply_gradients(grads_and_vars, global_step=global_step), []
225 
226   model_step = tf.Variable(0, trainable=False, collections=[], dtype=tf.int64)
227   accum_grads = []
228   accum_grads_and_vars = []
229   for grad, var in grads_and_vars:
230     accum_grad = tf.Variable(
231         tf.zeros(var.shape, dtype=grad.dtype),
232         trainable=False,
233         collections=[])
234     accum_grads.append(accum_grad)
235     accum_grads_and_vars.append((accum_grad, var))
236 
237   def _accum_grads(accum_fn=tf.assign_add, apply_gradients=False):
238     update_ops = [model_step.assign_add(1)]
239     for accum_grad, (grad, _) in zip(accum_grads, grads_and_vars):
240       with tf.control_dependencies([grad]):
241         update_ops.append(accum_fn(accum_grad, grad))
242     with tf.control_dependencies(update_ops):
243       if apply_gradients:
244         # Override the current name scope to create the optimizer slot variables
245         # in the same scope as if the optimizer was called outside of tf.cond.
246         # This is needed to ensure we can continue from a model trained without
247         # gradient accumulation (and vice-versa).
248         with tf.name_scope("%s/" % tf.get_variable_scope().name):
249           return _apply_gradients(accum_grads_and_vars, global_step=global_step)
250       else:
251         return tf.no_op()
252 
253   update_op = tf.cond(
254       tf.equal((model_step + 1) % accum_count, 0),
255       true_fn=lambda: _accum_grads(apply_gradients=True),
256       false_fn=lambda: tf.cond(
257           tf.equal(model_step % accum_count, 0),
258           true_fn=lambda: _accum_grads(accum_fn=tf.assign),
259           false_fn=_accum_grads))
260   extra_variables = accum_grads + [model_step]
261   return update_op, extra_variables
262 
263 def regularization_penalty(regularization_type, scale, weights_list=None):
264   """Computes the weights regularization penalty.
265 
266   Args:
267     regularization_type: The regularization type: ``l1``, ``l2``, or ``l1_l2``.
268     scale: The regularization multiplier. If :obj:`regularization_type` is
269       ``l1_l2``, this should be a list or tuple containing the L1 regularization
270       scale and the L2 regularization scale.
271     weights_list: The list of weights. Defaults to non bias variables.
272 
273   Returns:
274     The regularization penalty.
275 
276   Raises:
277     ValueError: if :obj:`regularization_type` is invalid or is ``l1_l2`` but
278       :obj:`scale` is not a sequence.
279   """
280   if weights_list is None:
281     weights_list = tf.trainable_variables()
282   weights_list = list(filter(lambda v: not _is_bias(v), weights_list))
283 
284   regularization_type = regularization_type.lower()
285   if regularization_type == "l1":
286     regularizer = tf.keras.regularizers.l1(l=float(scale))
287   elif regularization_type == "l2":
288     regularizer = tf.keras.regularizers.l2(l=float(scale))
289   elif regularization_type == "l1_l2":
290     if not isinstance(scale, collections.Sequence) or len(scale) != 2:
291       raise ValueError("l1_l2 regularization requires 2 scale values")
292     regularizer = tf.keras.regularizers.l1_l2(
293         l1=float(scale[0]), l2=float(scale[1]))
294   else:
295     raise ValueError("invalid regularization type %s" % regularization_type)
296 
297   penalty = tf.add_n([regularizer(w) for w in weights_list])
298   return penalty
299 
300 def _is_bias(variable):
301   return len(variable.shape.as_list()) == 1 and variable.name.endswith("bias:0")
302 
303 def _clip_gradients_by_norm(grads_and_vars, clip_gradients):
304   """Clips gradients by global norm."""
305   gradients, variables = zip(*grads_and_vars)
306   clipped_gradients, _ = tf.clip_by_global_norm(gradients, clip_gradients)
307   return list(zip(clipped_gradients, variables))
308 
309 def _summarize_gradients_norm(name, gradients):
310   """Summarizes global norm of gradients."""
311   tf.summary.scalar(name, tf.global_norm(list(zip(*gradients))[0]))
312
```
--- 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/optim.py b/opennmt/utils/optim.py
--- a/opennmt/utils/optim.py
+++ b/opennmt/utils/optim.py
@@ -1,6 +1,7 @@
 """Optimization related functions."""
 
 import collections
+import re
 
 import tensorflow as tf
 
@@ -183,6 +184,8 @@
       optimizer = DistributedOptimizer.from_params(optimizer, params=params.get("horovod"))
 
     # Gradients.
+    var_list = _get_trainable_variables(
+        var_list=var_list, freeze_variables=params.get("freeze_variables"))
     gradients = optimizer.compute_gradients(
         loss, var_list=var_list, colocate_gradients_with_ops=True)
     _summarize_gradients_norm("global_norm/gradient_norm", gradients)
@@ -309,3 +312,28 @@
 def _summarize_gradients_norm(name, gradients):
   """Summarizes global norm of gradients."""
   tf.summary.scalar(name, tf.global_norm(list(zip(*gradients))[0]))
+
+def _print_var_list(var_list):
+  for variable in var_list:
+    tf.logging.info(" * %s", variable.name)
+
+def _get_trainable_variables(var_list=None, freeze_variables=None):
+  if var_list is None:
+    var_list = tf.trainable_variables()
+  if freeze_variables:
+    if not isinstance(freeze_variables, list):
+      freeze_variables = [freeze_variables]
+    regexs = list(map(re.compile, freeze_variables))
+    frozen_variables = []
+    trainable_variables = []
+    for variable in var_list:
+      if any(regex.match(variable.name) for regex in regexs):
+        frozen_variables.append(variable)
+      else:
+        trainable_variables.append(variable)
+    tf.logging.info("Frozen variables:")
+    _print_var_list(frozen_variables)
+    tf.logging.info("Trainable variables:")
+    _print_var_list(trainable_variables)
+    var_list = trainable_variables
+  return var_list

{"golden_diff": "diff --git a/opennmt/utils/optim.py b/opennmt/utils/optim.py\n--- a/opennmt/utils/optim.py\n+++ b/opennmt/utils/optim.py\n@@ -1,6 +1,7 @@\n \"\"\"Optimization related functions.\"\"\"\n \n import collections\n+import re\n \n import tensorflow as tf\n \n@@ -183,6 +184,8 @@\n       optimizer = DistributedOptimizer.from_params(optimizer, params=params.get(\"horovod\"))\n \n     # Gradients.\n+    var_list = _get_trainable_variables(\n+        var_list=var_list, freeze_variables=params.get(\"freeze_variables\"))\n     gradients = optimizer.compute_gradients(\n         loss, var_list=var_list, colocate_gradients_with_ops=True)\n     _summarize_gradients_norm(\"global_norm/gradient_norm\", gradients)\n@@ -309,3 +312,28 @@\n def _summarize_gradients_norm(name, gradients):\n   \"\"\"Summarizes global norm of gradients.\"\"\"\n   tf.summary.scalar(name, tf.global_norm(list(zip(*gradients))[0]))\n+\n+def _print_var_list(var_list):\n+  for variable in var_list:\n+    tf.logging.info(\" * %s\", variable.name)\n+\n+def _get_trainable_variables(var_list=None, freeze_variables=None):\n+  if var_list is None:\n+    var_list = tf.trainable_variables()\n+  if freeze_variables:\n+    if not isinstance(freeze_variables, list):\n+      freeze_variables = [freeze_variables]\n+    regexs = list(map(re.compile, freeze_variables))\n+    frozen_variables = []\n+    trainable_variables = []\n+    for variable in var_list:\n+      if any(regex.match(variable.name) for regex in regexs):\n+        frozen_variables.append(variable)\n+      else:\n+        trainable_variables.append(variable)\n+    tf.logging.info(\"Frozen variables:\")\n+    _print_var_list(frozen_variables)\n+    tf.logging.info(\"Trainable variables:\")\n+    _print_var_list(trainable_variables)\n+    var_list = trainable_variables\n+  return var_list\n", "issue": "Fine-tuning only a subset of parameters\nHi,\r\n\r\nIs there the possibility of optimizing only a subset of parameters in the computation graph?\n", "before_files": [{"content": "\"\"\"Optimization related functions.\"\"\"\n\nimport collections\n\nimport tensorflow as tf\n\nfrom opennmt import optimizers\nfrom opennmt.utils import decay\nfrom opennmt.utils import misc\n\n\ndef learning_rate_decay_fn(decay_type,\n                           decay_rate,\n                           decay_steps,\n                           decay_step_duration=1,\n                           staircase=True,\n                           start_decay_steps=0,\n                           minimum_learning_rate=0):\n  \"\"\"Returns the learning rate decay functions.\n\n  Args:\n    decay_type: The type of decay. A function from ``tf.train`` or\n     :mod:`opennmt.utils.decay` as a string.\n    decay_rate: The decay rate to apply.\n    decay_steps: The decay steps as described in the decay type function.\n    decay_step_duration: The number of training steps that make 1 decay step.\n    staircase: If ``True``, learning rate is decayed in a staircase fashion.\n    start_decay_steps: Start decay after this many steps.\n    minimum_learning_rate: Do not decay past this learning rate value.\n\n  Returns:\n    A function with signature\n    ``(learning_rate, global_step) -> decayed_learning_rate``.\n\n  Raises:\n    ValueError: if :obj:`decay_type` can not be resolved.\n  \"\"\"\n  decay_params = {\n      \"decay_rate\": decay_rate,\n      \"decay_steps\": decay_steps,\n      \"staircase\": staircase\n  }\n  return learning_rate_decay_fn_v2(\n      decay_type,\n      decay_params=decay_params,\n      decay_step_duration=decay_step_duration,\n      start_decay_step=start_decay_steps,\n      minimum_learning_rate=minimum_learning_rate)\n\ndef learning_rate_decay_fn_v2(decay_type,\n                              decay_params=None,\n                              decay_step_duration=1,\n                              start_decay_step=0,\n                              minimum_learning_rate=0.0):\n  \"\"\"Returns the learning rate decay function.\n\n  Args:\n    decay_type: The type of decay. A function from ``tf.train`` or\n      :mod:`opennmt.utils.decay` as a string.\n    decay_params: Additional parameters for the decay function.\n    decay_step_duration: The number of training steps that make 1 decay step.\n    start_decay_step: Start decay after this many steps.\n    minimum_learning_rate: Do not decay past this learning rate value.\n\n  Returns:\n    A function with signature\n    ``(learning_rate, global_step) -> decayed_learning_rate``.\n\n  Raises:\n    ValueError: if :obj:`decay_type` can not be resolved.\n  \"\"\"\n  if decay_params is None:\n    decay_params = {}\n\n  def _decay_fn(learning_rate, global_step):\n    decay_op_name = None\n\n    if decay_op_name is None:\n      decay_op_name = getattr(tf.train, decay_type, None)\n    if decay_op_name is None:\n      decay_op_name = getattr(decay, decay_type, None)\n    if decay_op_name is None:\n      raise ValueError(\"Unknown decay function: {}\".format(decay_type))\n\n    # Map the training step to a decay step.\n    step = tf.maximum(global_step - start_decay_step, 0)\n    step //= decay_step_duration\n\n    learning_rate = decay_op_name(learning_rate, step, **decay_params)\n    return tf.maximum(learning_rate, minimum_learning_rate)\n\n  return _decay_fn\n\ndef get_optimizer_class(classname):\n  \"\"\"Returns the optimizer class.\n\n  Args:\n    classname: The name of the optimizer class in ``tf.train``,\n      ``tf.contrib.opt``, or ``opennmt.optimizers`` as a string.\n\n  Returns:\n    A class inheriting from ``tf.train.Optimizer``.\n\n  Raises:\n    ValueError: if :obj:`classname` can not be resolved.\n  \"\"\"\n  optimizer_class = None\n\n  if optimizer_class is None:\n    optimizer_class = getattr(tf.train, classname, None)\n  if optimizer_class is None:\n    optimizer_class = getattr(tf.contrib.opt, classname, None)\n  if optimizer_class is None:\n    optimizer_class = getattr(optimizers, classname, None)\n  if optimizer_class is None:\n    raise ValueError(\"Unknown optimizer class: {}\".format(classname))\n\n  return optimizer_class\n\ndef optimize(*args, **kwargs):\n  \"\"\"Wrapper around ``optimize_loss`` for backward compatibility.\"\"\"\n  update_op, _ = optimize_loss(*args, **kwargs)\n  return update_op\n\ndef optimize_loss(loss, params, mixed_precision=False, var_list=None, hvd=None):\n  \"\"\"Minimizes the loss.\n\n  Args:\n    loss: The loss to minimize.\n    params: A dictionary of hyperparameters.\n    mixed_precision: If ``True``, wraps the optimizer to maintain a float32 copy\n      of the weights.\n    var_list: The variables to update.\n    hvd: Optional Horovod object.\n\n  Returns:\n    The loss minimization op and a list of internal variables to initialize.\n  \"\"\"\n  regularization = params.get(\"regularization\")\n  if regularization is not None:\n    loss += regularization_penalty(\n        regularization[\"type\"], regularization[\"scale\"], weights_list=var_list)\n\n  global_step = tf.train.get_or_create_global_step()\n  with tf.variable_scope(\"optim\"):\n    learning_rate = tf.constant(params[\"learning_rate\"], dtype=tf.float32)\n    if params.get(\"decay_type\") is not None:\n      decay_params = params.get(\"decay_params\", {})\n      if \"decay_rate\" in params:\n        # Backward compatibility: fill params from previous options.\n        decay_params[\"decay_rate\"] = params[\"decay_rate\"]\n        decay_params[\"decay_steps\"] = params[\"decay_steps\"]\n        decay_params[\"staircase\"] = params.get(\"staircase\", True)\n      decay_fn = learning_rate_decay_fn_v2(\n          params[\"decay_type\"],\n          decay_params=decay_params,\n          decay_step_duration=params.get(\"decay_step_duration\", 1),\n          start_decay_step=params.get(\"start_decay_steps\", 0),\n          minimum_learning_rate=params.get(\"minimum_learning_rate\", 0))\n      learning_rate = decay_fn(learning_rate, global_step)\n    tf.summary.scalar(\"learning_rate\", learning_rate)\n\n    # Optimizer.\n    optimizer_class = get_optimizer_class(params[\"optimizer\"])\n    optimizer_params = params.get(\"optimizer_params\", {})\n    if optimizer_class.__name__ == \"AdafactorOptimizer\":\n      optimizer = optimizers.get_adafactor_optimizer_from_params(\n          optimizer_class, optimizer_params, learning_rate=learning_rate)\n    else:\n      weight_decay = params.get(\"weight_decay\")\n      if weight_decay is not None:\n        optimizer_class = tf.contrib.opt.extend_with_decoupled_weight_decay(optimizer_class)\n        optimizer = optimizer_class(\n            weight_decay, learning_rate=learning_rate, **optimizer_params)\n      else:\n        optimizer = optimizer_class(learning_rate, **optimizer_params)\n    if mixed_precision:\n      from opennmt.optimizers.mixed_precision_wrapper import get_loss_scale_from_params\n      optimizer = optimizers.MixedPrecisionOptimizerWrapper(\n          optimizer, loss_scale=get_loss_scale_from_params(params))\n    if hvd is not None:\n      from opennmt.optimizers.distributed_optimizer import DistributedOptimizer\n      optimizer = DistributedOptimizer.from_params(optimizer, params=params.get(\"horovod\"))\n\n    # Gradients.\n    gradients = optimizer.compute_gradients(\n        loss, var_list=var_list, colocate_gradients_with_ops=True)\n    _summarize_gradients_norm(\"global_norm/gradient_norm\", gradients)\n    if params.get(\"clip_gradients\") is not None:\n      gradients = _clip_gradients_by_norm(gradients, float(params[\"clip_gradients\"]))\n      _summarize_gradients_norm(\"global_norm/clipped_gradient_norm\", gradients)\n\n    return delayed_update(\n        optimizer,\n        gradients,\n        global_step,\n        accum_count=params.get(\"gradients_accum\", 1))\n\ndef delayed_update(optimizer, grads_and_vars, global_step, accum_count=1):\n  \"\"\"Possibly delays the parameters update by first accumulating gradients.\n\n  Args:\n    optimizer: The optimizer.\n    grads_and_vars: List of (gradient, variable) pairs.\n    global_step: The training step that will be increased when the parameters\n      are updated.\n    accum_count: The number of times to accumulate gradients, as a constant or\n      a ``tf.Tensor``.\n\n  Returns:\n    An operation that conditionally applies the gradients and a list of internal\n    variables to initialize.\n  \"\"\"\n\n  def _apply_gradients(grads_and_vars, global_step=None):\n    if hasattr(optimizer, \"allreduce_gradients\"):\n      grads_and_vars = optimizer.allreduce_gradients(grads_and_vars)\n    kwargs = {}\n    if \"decay_var_list\" in misc.function_args(optimizer.apply_gradients):\n      kwargs[\"decay_var_list\"] = [var for _, var in grads_and_vars if not _is_bias(var)]\n    return optimizer.apply_gradients(grads_and_vars, global_step=global_step, **kwargs)\n\n  if not tf.contrib.framework.is_tensor(accum_count) and accum_count == 1:\n    return _apply_gradients(grads_and_vars, global_step=global_step), []\n\n  model_step = tf.Variable(0, trainable=False, collections=[], dtype=tf.int64)\n  accum_grads = []\n  accum_grads_and_vars = []\n  for grad, var in grads_and_vars:\n    accum_grad = tf.Variable(\n        tf.zeros(var.shape, dtype=grad.dtype),\n        trainable=False,\n        collections=[])\n    accum_grads.append(accum_grad)\n    accum_grads_and_vars.append((accum_grad, var))\n\n  def _accum_grads(accum_fn=tf.assign_add, apply_gradients=False):\n    update_ops = [model_step.assign_add(1)]\n    for accum_grad, (grad, _) in zip(accum_grads, grads_and_vars):\n      with tf.control_dependencies([grad]):\n        update_ops.append(accum_fn(accum_grad, grad))\n    with tf.control_dependencies(update_ops):\n      if apply_gradients:\n        # Override the current name scope to create the optimizer slot variables\n        # in the same scope as if the optimizer was called outside of tf.cond.\n        # This is needed to ensure we can continue from a model trained without\n        # gradient accumulation (and vice-versa).\n        with tf.name_scope(\"%s/\" % tf.get_variable_scope().name):\n          return _apply_gradients(accum_grads_and_vars, global_step=global_step)\n      else:\n        return tf.no_op()\n\n  update_op = tf.cond(\n      tf.equal((model_step + 1) % accum_count, 0),\n      true_fn=lambda: _accum_grads(apply_gradients=True),\n      false_fn=lambda: tf.cond(\n          tf.equal(model_step % accum_count, 0),\n          true_fn=lambda: _accum_grads(accum_fn=tf.assign),\n          false_fn=_accum_grads))\n  extra_variables = accum_grads + [model_step]\n  return update_op, extra_variables\n\ndef regularization_penalty(regularization_type, scale, weights_list=None):\n  \"\"\"Computes the weights regularization penalty.\n\n  Args:\n    regularization_type: The regularization type: ``l1``, ``l2``, or ``l1_l2``.\n    scale: The regularization multiplier. If :obj:`regularization_type` is\n      ``l1_l2``, this should be a list or tuple containing the L1 regularization\n      scale and the L2 regularization scale.\n    weights_list: The list of weights. Defaults to non bias variables.\n\n  Returns:\n    The regularization penalty.\n\n  Raises:\n    ValueError: if :obj:`regularization_type` is invalid or is ``l1_l2`` but\n      :obj:`scale` is not a sequence.\n  \"\"\"\n  if weights_list is None:\n    weights_list = tf.trainable_variables()\n  weights_list = list(filter(lambda v: not _is_bias(v), weights_list))\n\n  regularization_type = regularization_type.lower()\n  if regularization_type == \"l1\":\n    regularizer = tf.keras.regularizers.l1(l=float(scale))\n  elif regularization_type == \"l2\":\n    regularizer = tf.keras.regularizers.l2(l=float(scale))\n  elif regularization_type == \"l1_l2\":\n    if not isinstance(scale, collections.Sequence) or len(scale) != 2:\n      raise ValueError(\"l1_l2 regularization requires 2 scale values\")\n    regularizer = tf.keras.regularizers.l1_l2(\n        l1=float(scale[0]), l2=float(scale[1]))\n  else:\n    raise ValueError(\"invalid regularization type %s\" % regularization_type)\n\n  penalty = tf.add_n([regularizer(w) for w in weights_list])\n  return penalty\n\ndef _is_bias(variable):\n  return len(variable.shape.as_list()) == 1 and variable.name.endswith(\"bias:0\")\n\ndef _clip_gradients_by_norm(grads_and_vars, clip_gradients):\n  \"\"\"Clips gradients by global norm.\"\"\"\n  gradients, variables = zip(*grads_and_vars)\n  clipped_gradients, _ = tf.clip_by_global_norm(gradients, clip_gradients)\n  return list(zip(clipped_gradients, variables))\n\ndef _summarize_gradients_norm(name, gradients):\n  \"\"\"Summarizes global norm of gradients.\"\"\"\n  tf.summary.scalar(name, tf.global_norm(list(zip(*gradients))[0]))\n", "path": "opennmt/utils/optim.py"}], "after_files": [{"content": "\"\"\"Optimization related functions.\"\"\"\n\nimport collections\nimport re\n\nimport tensorflow as tf\n\nfrom opennmt import optimizers\nfrom opennmt.utils import decay\nfrom opennmt.utils import misc\n\n\ndef learning_rate_decay_fn(decay_type,\n                           decay_rate,\n                           decay_steps,\n                           decay_step_duration=1,\n                           staircase=True,\n                           start_decay_steps=0,\n                           minimum_learning_rate=0):\n  \"\"\"Returns the learning rate decay functions.\n\n  Args:\n    decay_type: The type of decay. A function from ``tf.train`` or\n     :mod:`opennmt.utils.decay` as a string.\n    decay_rate: The decay rate to apply.\n    decay_steps: The decay steps as described in the decay type function.\n    decay_step_duration: The number of training steps that make 1 decay step.\n    staircase: If ``True``, learning rate is decayed in a staircase fashion.\n    start_decay_steps: Start decay after this many steps.\n    minimum_learning_rate: Do not decay past this learning rate value.\n\n  Returns:\n    A function with signature\n    ``(learning_rate, global_step) -> decayed_learning_rate``.\n\n  Raises:\n    ValueError: if :obj:`decay_type` can not be resolved.\n  \"\"\"\n  decay_params = {\n      \"decay_rate\": decay_rate,\n      \"decay_steps\": decay_steps,\n      \"staircase\": staircase\n  }\n  return learning_rate_decay_fn_v2(\n      decay_type,\n      decay_params=decay_params,\n      decay_step_duration=decay_step_duration,\n      start_decay_step=start_decay_steps,\n      minimum_learning_rate=minimum_learning_rate)\n\ndef learning_rate_decay_fn_v2(decay_type,\n                              decay_params=None,\n                              decay_step_duration=1,\n                              start_decay_step=0,\n                              minimum_learning_rate=0.0):\n  \"\"\"Returns the learning rate decay function.\n\n  Args:\n    decay_type: The type of decay. A function from ``tf.train`` or\n      :mod:`opennmt.utils.decay` as a string.\n    decay_params: Additional parameters for the decay function.\n    decay_step_duration: The number of training steps that make 1 decay step.\n    start_decay_step: Start decay after this many steps.\n    minimum_learning_rate: Do not decay past this learning rate value.\n\n  Returns:\n    A function with signature\n    ``(learning_rate, global_step) -> decayed_learning_rate``.\n\n  Raises:\n    ValueError: if :obj:`decay_type` can not be resolved.\n  \"\"\"\n  if decay_params is None:\n    decay_params = {}\n\n  def _decay_fn(learning_rate, global_step):\n    decay_op_name = None\n\n    if decay_op_name is None:\n      decay_op_name = getattr(tf.train, decay_type, None)\n    if decay_op_name is None:\n      decay_op_name = getattr(decay, decay_type, None)\n    if decay_op_name is None:\n      raise ValueError(\"Unknown decay function: {}\".format(decay_type))\n\n    # Map the training step to a decay step.\n    step = tf.maximum(global_step - start_decay_step, 0)\n    step //= decay_step_duration\n\n    learning_rate = decay_op_name(learning_rate, step, **decay_params)\n    return tf.maximum(learning_rate, minimum_learning_rate)\n\n  return _decay_fn\n\ndef get_optimizer_class(classname):\n  \"\"\"Returns the optimizer class.\n\n  Args:\n    classname: The name of the optimizer class in ``tf.train``,\n      ``tf.contrib.opt``, or ``opennmt.optimizers`` as a string.\n\n  Returns:\n    A class inheriting from ``tf.train.Optimizer``.\n\n  Raises:\n    ValueError: if :obj:`classname` can not be resolved.\n  \"\"\"\n  optimizer_class = None\n\n  if optimizer_class is None:\n    optimizer_class = getattr(tf.train, classname, None)\n  if optimizer_class is None:\n    optimizer_class = getattr(tf.contrib.opt, classname, None)\n  if optimizer_class is None:\n    optimizer_class = getattr(optimizers, classname, None)\n  if optimizer_class is None:\n    raise ValueError(\"Unknown optimizer class: {}\".format(classname))\n\n  return optimizer_class\n\ndef optimize(*args, **kwargs):\n  \"\"\"Wrapper around ``optimize_loss`` for backward compatibility.\"\"\"\n  update_op, _ = optimize_loss(*args, **kwargs)\n  return update_op\n\ndef optimize_loss(loss, params, mixed_precision=False, var_list=None, hvd=None):\n  \"\"\"Minimizes the loss.\n\n  Args:\n    loss: The loss to minimize.\n    params: A dictionary of hyperparameters.\n    mixed_precision: If ``True``, wraps the optimizer to maintain a float32 copy\n      of the weights.\n    var_list: The variables to update.\n    hvd: Optional Horovod object.\n\n  Returns:\n    The loss minimization op and a list of internal variables to initialize.\n  \"\"\"\n  regularization = params.get(\"regularization\")\n  if regularization is not None:\n    loss += regularization_penalty(\n        regularization[\"type\"], regularization[\"scale\"], weights_list=var_list)\n\n  global_step = tf.train.get_or_create_global_step()\n  with tf.variable_scope(\"optim\"):\n    learning_rate = tf.constant(params[\"learning_rate\"], dtype=tf.float32)\n    if params.get(\"decay_type\") is not None:\n      decay_params = params.get(\"decay_params\", {})\n      if \"decay_rate\" in params:\n        # Backward compatibility: fill params from previous options.\n        decay_params[\"decay_rate\"] = params[\"decay_rate\"]\n        decay_params[\"decay_steps\"] = params[\"decay_steps\"]\n        decay_params[\"staircase\"] = params.get(\"staircase\", True)\n      decay_fn = learning_rate_decay_fn_v2(\n          params[\"decay_type\"],\n          decay_params=decay_params,\n          decay_step_duration=params.get(\"decay_step_duration\", 1),\n          start_decay_step=params.get(\"start_decay_steps\", 0),\n          minimum_learning_rate=params.get(\"minimum_learning_rate\", 0))\n      learning_rate = decay_fn(learning_rate, global_step)\n    tf.summary.scalar(\"learning_rate\", learning_rate)\n\n    # Optimizer.\n    optimizer_class = get_optimizer_class(params[\"optimizer\"])\n    optimizer_params = params.get(\"optimizer_params\", {})\n    if optimizer_class.__name__ == \"AdafactorOptimizer\":\n      optimizer = optimizers.get_adafactor_optimizer_from_params(\n          optimizer_class, optimizer_params, learning_rate=learning_rate)\n    else:\n      weight_decay = params.get(\"weight_decay\")\n      if weight_decay is not None:\n        optimizer_class = tf.contrib.opt.extend_with_decoupled_weight_decay(optimizer_class)\n        optimizer = optimizer_class(\n            weight_decay, learning_rate=learning_rate, **optimizer_params)\n      else:\n        optimizer = optimizer_class(learning_rate, **optimizer_params)\n    if mixed_precision:\n      from opennmt.optimizers.mixed_precision_wrapper import get_loss_scale_from_params\n      optimizer = optimizers.MixedPrecisionOptimizerWrapper(\n          optimizer, loss_scale=get_loss_scale_from_params(params))\n    if hvd is not None:\n      from opennmt.optimizers.distributed_optimizer import DistributedOptimizer\n      optimizer = DistributedOptimizer.from_params(optimizer, params=params.get(\"horovod\"))\n\n    # Gradients.\n    var_list = _get_trainable_variables(\n        var_list=var_list, freeze_variables=params.get(\"freeze_variables\"))\n    gradients = optimizer.compute_gradients(\n        loss, var_list=var_list, colocate_gradients_with_ops=True)\n    _summarize_gradients_norm(\"global_norm/gradient_norm\", gradients)\n    if params.get(\"clip_gradients\") is not None:\n      gradients = _clip_gradients_by_norm(gradients, float(params[\"clip_gradients\"]))\n      _summarize_gradients_norm(\"global_norm/clipped_gradient_norm\", gradients)\n\n    return delayed_update(\n        optimizer,\n        gradients,\n        global_step,\n        accum_count=params.get(\"gradients_accum\", 1))\n\ndef delayed_update(optimizer, grads_and_vars, global_step, accum_count=1):\n  \"\"\"Possibly delays the parameters update by first accumulating gradients.\n\n  Args:\n    optimizer: The optimizer.\n    grads_and_vars: List of (gradient, variable) pairs.\n    global_step: The training step that will be increased when the parameters\n      are updated.\n    accum_count: The number of times to accumulate gradients, as a constant or\n      a ``tf.Tensor``.\n\n  Returns:\n    An operation that conditionally applies the gradients and a list of internal\n    variables to initialize.\n  \"\"\"\n\n  def _apply_gradients(grads_and_vars, global_step=None):\n    if hasattr(optimizer, \"allreduce_gradients\"):\n      grads_and_vars = optimizer.allreduce_gradients(grads_and_vars)\n    kwargs = {}\n    if \"decay_var_list\" in misc.function_args(optimizer.apply_gradients):\n      kwargs[\"decay_var_list\"] = [var for _, var in grads_and_vars if not _is_bias(var)]\n    return optimizer.apply_gradients(grads_and_vars, global_step=global_step, **kwargs)\n\n  if not tf.contrib.framework.is_tensor(accum_count) and accum_count == 1:\n    return _apply_gradients(grads_and_vars, global_step=global_step), []\n\n  model_step = tf.Variable(0, trainable=False, collections=[], dtype=tf.int64)\n  accum_grads = []\n  accum_grads_and_vars = []\n  for grad, var in grads_and_vars:\n    accum_grad = tf.Variable(\n        tf.zeros(var.shape, dtype=grad.dtype),\n        trainable=False,\n        collections=[])\n    accum_grads.append(accum_grad)\n    accum_grads_and_vars.append((accum_grad, var))\n\n  def _accum_grads(accum_fn=tf.assign_add, apply_gradients=False):\n    update_ops = [model_step.assign_add(1)]\n    for accum_grad, (grad, _) in zip(accum_grads, grads_and_vars):\n      with tf.control_dependencies([grad]):\n        update_ops.append(accum_fn(accum_grad, grad))\n    with tf.control_dependencies(update_ops):\n      if apply_gradients:\n        # Override the current name scope to create the optimizer slot variables\n        # in the same scope as if the optimizer was called outside of tf.cond.\n        # This is needed to ensure we can continue from a model trained without\n        # gradient accumulation (and vice-versa).\n        with tf.name_scope(\"%s/\" % tf.get_variable_scope().name):\n          return _apply_gradients(accum_grads_and_vars, global_step=global_step)\n      else:\n        return tf.no_op()\n\n  update_op = tf.cond(\n      tf.equal((model_step + 1) % accum_count, 0),\n      true_fn=lambda: _accum_grads(apply_gradients=True),\n      false_fn=lambda: tf.cond(\n          tf.equal(model_step % accum_count, 0),\n          true_fn=lambda: _accum_grads(accum_fn=tf.assign),\n          false_fn=_accum_grads))\n  extra_variables = accum_grads + [model_step]\n  return update_op, extra_variables\n\ndef regularization_penalty(regularization_type, scale, weights_list=None):\n  \"\"\"Computes the weights regularization penalty.\n\n  Args:\n    regularization_type: The regularization type: ``l1``, ``l2``, or ``l1_l2``.\n    scale: The regularization multiplier. If :obj:`regularization_type` is\n      ``l1_l2``, this should be a list or tuple containing the L1 regularization\n      scale and the L2 regularization scale.\n    weights_list: The list of weights. Defaults to non bias variables.\n\n  Returns:\n    The regularization penalty.\n\n  Raises:\n    ValueError: if :obj:`regularization_type` is invalid or is ``l1_l2`` but\n      :obj:`scale` is not a sequence.\n  \"\"\"\n  if weights_list is None:\n    weights_list = tf.trainable_variables()\n  weights_list = list(filter(lambda v: not _is_bias(v), weights_list))\n\n  regularization_type = regularization_type.lower()\n  if regularization_type == \"l1\":\n    regularizer = tf.keras.regularizers.l1(l=float(scale))\n  elif regularization_type == \"l2\":\n    regularizer = tf.keras.regularizers.l2(l=float(scale))\n  elif regularization_type == \"l1_l2\":\n    if not isinstance(scale, collections.Sequence) or len(scale) != 2:\n      raise ValueError(\"l1_l2 regularization requires 2 scale values\")\n    regularizer = tf.keras.regularizers.l1_l2(\n        l1=float(scale[0]), l2=float(scale[1]))\n  else:\n    raise ValueError(\"invalid regularization type %s\" % regularization_type)\n\n  penalty = tf.add_n([regularizer(w) for w in weights_list])\n  return penalty\n\ndef _is_bias(variable):\n  return len(variable.shape.as_list()) == 1 and variable.name.endswith(\"bias:0\")\n\ndef _clip_gradients_by_norm(grads_and_vars, clip_gradients):\n  \"\"\"Clips gradients by global norm.\"\"\"\n  gradients, variables = zip(*grads_and_vars)\n  clipped_gradients, _ = tf.clip_by_global_norm(gradients, clip_gradients)\n  return list(zip(clipped_gradients, variables))\n\ndef _summarize_gradients_norm(name, gradients):\n  \"\"\"Summarizes global norm of gradients.\"\"\"\n  tf.summary.scalar(name, tf.global_norm(list(zip(*gradients))[0]))\n\ndef _print_var_list(var_list):\n  for variable in var_list:\n    tf.logging.info(\" * %s\", variable.name)\n\ndef _get_trainable_variables(var_list=None, freeze_variables=None):\n  if var_list is None:\n    var_list = tf.trainable_variables()\n  if freeze_variables:\n    if not isinstance(freeze_variables, list):\n      freeze_variables = [freeze_variables]\n    regexs = list(map(re.compile, freeze_variables))\n    frozen_variables = []\n    trainable_variables = []\n    for variable in var_list:\n      if any(regex.match(variable.name) for regex in regexs):\n        frozen_variables.append(variable)\n      else:\n        trainable_variables.append(variable)\n    tf.logging.info(\"Frozen variables:\")\n    _print_var_list(frozen_variables)\n    tf.logging.info(\"Trainable variables:\")\n    _print_var_list(trainable_variables)\n    var_list = trainable_variables\n  return var_list\n", "path": "opennmt/utils/optim.py"}]}