Repository: rwth-i6/returnn

Issue:
`decouple_constraints` more efficient and more generic
`decouple_constraints` currently is implemented in a very simple way. It first applies the whole standard update for all variables and then it applies another update for the constraints. However, this second update waits for all the first updates to finish, and this is not per parameter but for the whole model. Specifically:
```python
with tf_compat.v1.variable_scope("optimize_constraints"):
  with tf_compat.v1.variable_scope("factor"):
    factor = (self.get_current_step_learning_rate() / float(self.initial_learning_rate))
    factor *= self.config.float("decouple_constraints_factor", 0.025)
  sgd_optimizer = tf_compat.v1.train.GradientDescentOptimizer(
    learning_rate=factor, use_locking=self.use_locking)
  with tf.control_dependencies([self.optim_op]):
    self.optim_op = sgd_optimizer.minimize(self.constraints, var_list=self.trainable_vars)
```
`optim_op` before this was the single grouped op with all the updates.

This can be more efficient by using the `tf.Variable` `constraint` option which is specifically made for this purpose, and basically is the same but more efficient, per parameter:
```python
class _DenseResourceVariableProcessor(_OptimizableVariable):
  ...
  def update_op(self, optimizer, g):
    ...
    update_op = optimizer._resource_apply_dense(g, self._v)
    if self._v.constraint is not None:
      with ops.control_dependencies([update_op]):
        return self._v.assign(self._v.constraint(self._v))
    else:
      return update_op
```

Further, we could extend RETURNN to allow for other custom `constraint`.


Discussion:
Note that the TFA implementation of decoupled weight decay does actually the reverse order, i.e. it first applies weight decay, and then the gradient update. See [here](https://github.com/tensorflow/addons/blob/8cec33fcaaf1cf90aec7bdd55a0fcdbb251ce5c2/tensorflow_addons/optimizers/weight_decay_optimizers.py#L24). This should be more efficient. However, there is a race condition which leads to slightly non-deterministic behavior: The weight decay op does not have any control dependencies, so the time it is applied is arbitrary (only it must be before the grad update). But for calculating the forward + backprop, it is arbitrary whether this uses the weight decayed var or the original var. (I reported that here: https://github.com/tensorflow/addons/issues/2683)

Resolving patch:
diff --git a/returnn/tf/updater.py b/returnn/tf/updater.py
--- a/returnn/tf/updater.py
+++ b/returnn/tf/updater.py
@@ -7,6 +7,7 @@
 from __future__ import print_function
 
 import typing
+from collections import OrderedDict
 import tensorflow as tf
 from tensorflow.python.ops import resource_variable_ops
 
@@ -117,6 +118,8 @@ class Updater(object):
   and the update-op for all trainable vars.
   See the code of :func:`Updater.create_optimizer` for valid config options.
 
+  Wraps one or multiple tf.compat.v1.train.Optimizer, and extends it by some further functionality.
+
   Note: `Vincent Vanhoucke says <https://github.com/tensorflow/tensorflow/issues/323#issuecomment-159116515>`_,
   in case you get nans, consider increasing the epsilon (for Adam, Nadam and similar).
   This is the config option ``optimizer_epsilon``.
@@ -154,24 +157,26 @@ class Updater(object):
 
   def __init__(self, config, network, initial_learning_rate=1.):
     """
-    :param Config.Config config:
+    :param returnn.config.Config config:
     :param TFNetwork network:
     :param float initial_learning_rate:
     """
     self.config = config
     self.learning_rate_var = tf.Variable(name="learning_rate", initial_value=0.0, trainable=False, dtype="float32")
+    self.initial_learning_rate = initial_learning_rate
+    self.learning_rate = None  # type: typing.Optional[tf.Tensor]
     self.trainable_vars = []  # type: typing.List[tf.Variable]
     self.network = network
+    self.global_train_step = self.network.global_train_step
     self.use_locking = self.config.bool("optimizer_use_locking", False)
-    self.initial_learning_rate = initial_learning_rate
     if self.config.bool("decouple_constraints", False):
       # https://arxiv.org/abs/1711.05101, Fixing Weight Decay Regularization in Adam
       self.loss = network.get_total_loss()
-      self.constraints = network.get_total_constraints()
+      self.decoupled_constraints = network.get_total_constraints()
     else:
       self.loss = network.get_objective()
-      self.constraints = None
-    self.optimizer = None  # type: typing.Optional[WrapOptimizer]
+      self.decoupled_constraints = None
+    self.optimizers = OrderedDict()  # optimizer_opts|None -> tf.compat.v1.train.Optimizer
     self.optim_op = None  # type: typing.Optional[tf.Operation]
     self.optim_meta_losses_dict = None  # type: typing.Optional[typing.Dict[str,tf.Tensor]]
     self.optimizer_vars = []  # type: typing.List[tf.Variable]
@@ -263,23 +268,19 @@ def create_optim_op(self):
         trainable_vars_custom_update.append(v)
         trainable_vars_for_gradients.remove(v)
 
-    if not self.optimizer:
-      self.optimizer = WrapOptimizer(
-        config=self.config,
-        learning_rate=self.get_current_step_learning_rate(),
-        global_train_step=self.network.global_train_step,
-        use_locking=self.use_locking)
-      self.optimizer.create_all_needed_optimizers(trainable_vars_for_gradients)
+    self.learning_rate = self.get_current_step_learning_rate()
+    self.optimizers.clear()
+    self.create_all_needed_optimizers(trainable_vars_for_gradients)
 
     with tf_compat.v1.variable_scope("optimize"):
       meta_losses_scope = MetaLosses.enter_gradient_scope()
-      apply_grads = self.optimizer.get_apply_grads_op(self.loss, trainable_vars_for_gradients)
+      apply_grads = self.get_apply_grads_op(self.loss, trainable_vars_for_gradients)
       meta_losses_scope.exit()
       self.optim_meta_losses_dict = meta_losses_scope.losses_as_fetch_dict()
       if meta_losses_scope.losses:
         with tf.name_scope("meta_loss"):
           meta_loss = meta_losses_scope.summed_loss_for_optimization()
-          meta_apply_grads = self.optimizer.get_apply_grads_op(meta_loss, trainable_vars_for_gradients)
+          meta_apply_grads = self.get_apply_grads_op(meta_loss, trainable_vars_for_gradients)
         apply_grads = tf.group(apply_grads, meta_apply_grads)
       self.optim_op = apply_grads
 
@@ -292,16 +293,6 @@ def create_optim_op(self):
           updates.append(custom_update.update_var(param))
         self.optim_op = tf.group(*updates)
 
-    if self.constraints is not None:
-      with tf_compat.v1.variable_scope("optimize_constraints"):
-        with tf_compat.v1.variable_scope("factor"):
-          factor = (self.get_current_step_learning_rate() / float(self.initial_learning_rate))
-          factor *= self.config.float("decouple_constraints_factor", 0.025)
-        sgd_optimizer = tf_compat.v1.train.GradientDescentOptimizer(
-          learning_rate=factor, use_locking=self.use_locking)
-        with tf.control_dependencies([self.optim_op]):
-          self.optim_op = sgd_optimizer.minimize(self.constraints, var_list=self.trainable_vars)
-
     if self.config.opt_typed_value("extra_updates"):
       extra_updates = self.config.typed_dict["extra_updates"]
       assert isinstance(extra_updates, dict)  # dict var_name -> function(var)
@@ -329,13 +320,13 @@ def create_optim_op(self):
           extra_updates_op_list.append(op)
         self.optim_op = tf.group(self.optim_op, *extra_updates_op_list)
 
-    slot_names_per_optimizer = self.optimizer.get_slot_names_per_optimizer()
+    slot_names_per_optimizer = self.get_slot_names_per_optimizer()
     slot_vars = []
     for opt_key, slot_names in slot_names_per_optimizer.items():
       print("Initialize optimizer (%s) with slots %s." % (opt_key or "default", slot_names), file=log.v3)
       for slot_name in slot_names:
-        for v in self.optimizer.filter_var_list_per_optimizer_key(trainable_vars_for_gradients, opt_key=opt_key):
-          slot_var = self.optimizer.get_slot(var=v, name=slot_name)
+        for v in self.filter_var_list_per_optimizer_key(trainable_vars_for_gradients, opt_key=opt_key):
+          slot_var = self.get_slot(var=v, name=slot_name)
           if slot_var is None:
             print("Warning: No slot_var found for variable %r, slot_name %r. Maybe no gradient for this var?" % (
               v, slot_name), file=log.v3)
@@ -406,49 +397,6 @@ def init_optimizer_vars(self, session):
     self.get_optim_op()  # make sure it is initialized
     session.run(self.optimizer_init_vars_op)
 
-
-def accum_grad_multiple_step(grad, var, train_step, num_accum_steps):
-  """
-  :param tf.Tensor|tf.IndexedSlices grad:
-  :param tf.Variable var:
-  :param tf.Tensor train_step: int, scalar
-  :param int num_accum_steps:
-  :return: modified grad
-  :rtype: tf.Tensor
-  """
-  from returnn.tf.util.basic import reuse_name_scope_of_tensor, get_base_name
-  with reuse_name_scope_of_tensor(grad, postfix="/%s_accum_grad" % get_base_name(grad)):
-    shape = var.get_shape().as_list()
-    v = tf_compat.v1.get_variable(
-      name="var_accum_grad", shape=shape, dtype=grad.dtype,
-      initializer=tf.zeros_initializer(), trainable=False)
-    return tf.cond(
-      tf.less_equal(tf_compat.v1.mod(train_step, num_accum_steps), 0),
-      lambda: tf_compat.v1.assign(v, grad),
-      lambda: tf_compat.v1.assign_add(v, grad))
-
-
-class WrapOptimizer:
-  """
-  Wraps a tf.compat.v1.train.Optimizer (or multiple).
-  This is wrapped for a simpler interface, and also to allow for multiple optimizers.
-  This class is not derived from tf.compat.v1.train.Optimizer itself, to keep it simple.
-  """
-
-  def __init__(self, config, learning_rate, global_train_step, use_locking):
-    """
-    :param Config.Config config:
-    :param tf.Tensor learning_rate:
-    :param tf.Tensor global_train_step:
-    :param bool use_locking:
-    """
-    self.config = config
-    self.learning_rate = learning_rate
-    self.global_train_step = global_train_step
-    self.use_locking = use_locking
-    from collections import OrderedDict
-    self.optimizers = OrderedDict()  # optimizer_opts|None -> tf.compat.v1.train.Optimizer
-
   def get_default_optimizer(self):
     """
     :rtype: tf.compat.v1.train.Optimizer
@@ -684,13 +632,13 @@ def get_slot(self, var, name):
     return opt.get_slot(var, name)
 
   class _GetGlobalInfo:
-    def __init__(self, optimizer, all_vars, var_grads):
+    def __init__(self, updater, all_vars, var_grads):
       """
-      :param WrapOptimizer optimizer:
+      :param Updater updater:
       :param list[tf.Variable] all_vars:
       :param dict[tf.Variable,tf.Tensor] var_grads:
       """
-      self.optimizer = optimizer
+      self.updater = updater
       self.all_vars = all_vars
       self.var_grads = var_grads
       self.all_grads = list(var_grads.values())  # not necessarily the same length as all_vars
@@ -707,7 +655,7 @@ def _build_vars_by_tag_dict(self):
       """
       res = {}
       for var in self.all_vars:
-        opts = self.optimizer._get_updater_opts_from_var(var)
+        opts = self.updater._get_updater_opts_from_var(var)
         var_tags = opts.get("tags", [])
         for tag in var_tags:
           res.setdefault(tag, set()).add(var)
@@ -765,7 +713,7 @@ def get_global_grad_norm_per_tag(self, tag):
         from returnn.tf.util.basic import get_valid_scope_name_from_str
         with tf.name_scope("global_norm_for_tag_%s" % get_valid_scope_name_from_str(tag)):
           norm = self._global_norm({self.var_grads[var] for var in self.vars_by_tag[tag]})
-        if self.optimizer.config.bool_or_other("debug_grad_summaries", False):
+        if self.updater.config.bool_or_other("debug_grad_summaries", False):
           tf_compat.v1.summary.scalar("global_norm_for_tag_%s" % get_valid_scope_name_from_str(tag), norm)
         self._global_grad_norm_per_tag[tag] = norm
       return self._global_grad_norm_per_tag[tag]
@@ -778,7 +726,7 @@ def get_maximize_grad_norm_var_grads(self, factor):
       """
       if self._maximize_grad_norm_var_grads is None:
         loss_ext = self.get_global_grad_norm() * (-factor)
-        grads_and_vars_ext = self.optimizer._compute_gradients(loss_ext, var_list=self.all_vars)
+        grads_and_vars_ext = self.updater._compute_gradients(loss_ext, var_list=self.all_vars)
         self._maximize_grad_norm_var_grads = {var: grad for (grad, var) in grads_and_vars_ext if grad is not None}
       return self._maximize_grad_norm_var_grads
 
@@ -936,7 +884,7 @@ def get_apply_grads_op(self, loss, var_list):
     var_grads = {var: grad for (grad, var) in grads_and_vars if grad is not None}
     if not var_grads:
       raise Exception("no single variable to train")
-    global_info = self._GetGlobalInfo(optimizer=self, all_vars=var_list, var_grads=var_grads)
+    global_info = self._GetGlobalInfo(updater=self, all_vars=var_list, var_grads=var_grads)
     if self.config.bool_or_other("debug_grad_summaries", False):
       tf_compat.v1.summary.scalar("global_grad_norm", global_info.get_global_grad_norm())
     grads_per_apply_grad_opts = {}  # dict apply_grad_opts -> list of (grad, var)
@@ -947,6 +895,34 @@ def get_apply_grads_op(self, loss, var_list):
       new_grad, apply_grad_opts = self._post_process_grad(grad=grad, var=var, global_info=global_info)
       grads_per_apply_grad_opts.setdefault(make_hashable(apply_grad_opts), []).append((new_grad, var))
 
+    if self.decoupled_constraints is not None:
+      # Note: We want to perform the decoupled constraint updates after all the gradients (and post processing)
+      # is calculated (i.e. forward + backprop used the original variable, not any weight decayed version).
+      # We want to perform the decoupled constraint updates before we do the gradient update.
+      # This is consistent to other frameworks, e.g. PyTorch.
+      # https://github.com/rwth-i6/returnn/issues/1007
+      # The constraints are given as losses (e.g. L2 norm of the weights) thus we use SGD
+      # which is then equivalent to the standard weight decay.
+      # Also see the paper: https://arxiv.org/abs/1711.05101, Fixing Weight Decay Regularization in Adam
+      with tf_compat.v1.variable_scope("optimize_constraints"):
+        with tf_compat.v1.variable_scope("factor"):
+          factor = (self.learning_rate / float(self.initial_learning_rate))
+          factor *= self.config.float("decouple_constraints_factor", 0.025)
+        constraint_grad_descent_optimizer = tf_compat.v1.train.GradientDescentOptimizer(
+          learning_rate=factor, use_locking=self.use_locking)
+        for _, grads_and_vars_per_opts in grads_per_apply_grad_opts.items():
+          for i, (grad, var) in enumerate(grads_and_vars_per_opts):
+            (constraint_grad, var_), = constraint_grad_descent_optimizer.compute_gradients(
+              self.decoupled_constraints, var_list=[var])
+            assert var_ is var
+            if constraint_grad is None:  # no constraint on this var
+              continue
+            with tf.control_dependencies([grad]):
+              apply_constraint = constraint_grad_descent_optimizer.apply_gradients([(constraint_grad, var)])
+            with tf.control_dependencies([apply_constraint]):
+              grad = tf.identity(grad)
+            grads_and_vars_per_opts[i] = (grad, var)
+
     all_apply_grads = []
     assert grads_per_apply_grad_opts
     for apply_grad_opts, grads_and_vars_per_opts in grads_per_apply_grad_opts.items():
@@ -956,6 +932,27 @@ def get_apply_grads_op(self, loss, var_list):
     return tf.group(*all_apply_grads)
 
 
+def accum_grad_multiple_step(grad, var, train_step, num_accum_steps):
+  """
+  :param tf.Tensor|tf.IndexedSlices grad:
+  :param tf.Variable var:
+  :param tf.Tensor train_step: int, scalar
+  :param int num_accum_steps:
+  :return: modified grad
+  :rtype: tf.Tensor
+  """
+  from returnn.tf.util.basic import reuse_name_scope_of_tensor, get_base_name
+  with reuse_name_scope_of_tensor(grad, postfix="/%s_accum_grad" % get_base_name(grad)):
+    shape = var.get_shape().as_list()
+    v = tf_compat.v1.get_variable(
+      name="var_accum_grad", shape=shape, dtype=grad.dtype,
+      initializer=tf.zeros_initializer(), trainable=False)
+    return tf.cond(
+      tf.less_equal(tf_compat.v1.mod(train_step, num_accum_steps), 0),
+      lambda: tf_compat.v1.assign(v, grad),
+      lambda: tf_compat.v1.assign_add(v, grad))
+
+
 class _KerasOptimizerWrapper(Optimizer):
   """
   Wraps a TF optimizer into a standard TF optimizer.


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -871,9 +871,9 @@ def find_op_by_type(type_name):
     updater.set_learning_rate(0.1, session=session)
     updater.init_optimizer_vars(session=session)
     optim_op = updater.get_optim_op()
-    assert isinstance(updater.optimizer.get_default_optimizer(), tf_compat.v1.train.AdamOptimizer)
-    adam_weights_m_t = updater.optimizer.get_slot(var=weights_t, name="m")
-    adam_weights_v_t = updater.optimizer.get_slot(var=weights_t, name="v")
+    assert isinstance(updater.get_default_optimizer(), tf_compat.v1.train.AdamOptimizer)
+    adam_weights_m_t = updater.get_slot(var=weights_t, name="m")
+    adam_weights_v_t = updater.get_slot(var=weights_t, name="v")
     assert isinstance(adam_weights_m_t, tf.Variable)
     assert isinstance(adam_weights_v_t, tf.Variable)
     summaries_t = tf_compat.v1.summary.merge_all()
diff --git a/tests/test_TFUpdater.py b/tests/test_TFUpdater.py
--- a/tests/test_TFUpdater.py
+++ b/tests/test_TFUpdater.py
@@ -233,6 +233,47 @@ def test_Updater_simple_batch():
     session.run(updater.get_optim_op(), feed_dict=feed_dict)
 
 
+def test_Updater_decouple_constraints():
+  with make_scope() as session:
+    from returnn.tf.network import TFNetwork, ExternData
+    from returnn.config import Config
+    from returnn.datasets.generating import Task12AXDataset
+    dataset = Task12AXDataset()
+    dataset.init_seq_order(epoch=1)
+    extern_data = ExternData()
+    extern_data.init_from_dataset(dataset)
+
+    config = Config({"decouple_constraints": True})
+    network = TFNetwork(extern_data=extern_data, train_flag=True)
+    network.construct_from_dict({
+      "layer1": {"class": "linear", "activation": "tanh", "n_out": 13, "L2": 0.01, "from": "data:data"},
+      "layer2": {"class": "linear", "activation": "tanh", "n_out": 13, "L2": 0.01, "from": "layer1"},
+      "output": {"class": "softmax", "loss": "ce", "target": "classes", "from": "layer2"}
+    })
+    network.initialize_params(session=session)
+
+    updater = Updater(config=config, network=network)
+    updater.set_learning_rate(1.0, session=session)
+    updater.set_trainable_vars(network.get_trainable_params())
+    updater.init_optimizer_vars(session=session)
+    update_op = updater.get_optim_op()
+    assert updater.decoupled_constraints is not None
+
+    from returnn.tf.data_pipeline import FeedDictDataProvider
+    batches = dataset.generate_batches(
+      recurrent_net=network.recurrent,
+      batch_size=100,
+      max_seqs=10,
+      max_seq_length=sys.maxsize,
+      used_data_keys=network.used_data_keys)
+    data_provider = FeedDictDataProvider(
+      tf_session=session, extern_data=extern_data,
+      data_keys=network.used_data_keys,
+      dataset=dataset, batches=batches)
+    feed_dict, _ = data_provider.get_feed_dict(single_threaded=True)
+    session.run(update_op, feed_dict=feed_dict)
+
+
 def test_Updater_multiple_optimizers():
   with make_scope() as session:
     from returnn.tf.network import TFNetwork, ExternData
@@ -260,8 +301,7 @@ def test_Updater_multiple_optimizers():
     updater.init_optimizer_vars(session=session)
 
     optim_op = updater.get_optim_op()
-    assert isinstance(updater.optimizer, WrapOptimizer)
-    assert len(updater.optimizer.optimizers) == 3
+    assert len(updater.optimizers) == 3
 
     from returnn.tf.data_pipeline import FeedDictDataProvider
     batches = dataset.generate_batches(
@@ -306,8 +346,7 @@ def test_Updater_multiple_optimizers_and_opts():
     updater.init_optimizer_vars(session=session)
 
     optim_op = updater.get_optim_op()
-    assert isinstance(updater.optimizer, WrapOptimizer)
-    assert len(updater.optimizer.optimizers) == 3
+    assert len(updater.optimizers) == 3
 
     from returnn.tf.data_pipeline import FeedDictDataProvider
     batches = dataset.generate_batches(