Repository: OpenNMT/OpenNMT-tf
Pull request #736 (2020-11-05 09:58:02)

Cleanup moving average for mirrored variables
There is no need to get the primary variable to make it work. The update should be run under the distributed strategy scope.


Patch:
diff --git a/opennmt/training.py b/opennmt/training.py
index ba9a4de1..6399382a 100644
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -7,7 +7,6 @@ import time
 import tensorflow as tf
 
 from opennmt.optimizers import utils as optimizer_util
-from opennmt.utils import misc
 
 
 class Trainer:
@@ -104,8 +103,7 @@ class Trainer:
                 self._model.trainable_variables,
                 iterations,
                 decay=moving_average_decay)
-          else:
-            moving_average.update()
+          self._update_moving_average(moving_average)
 
         step = iterations.numpy()
         reset_throughput = False
@@ -319,6 +317,10 @@ class Trainer:
         scale=self._gradient_accumulator.step)
     self._gradient_accumulator.reset()
 
+  def _update_moving_average(self, moving_average):
+    """Updates the moving average of variables."""
+    moving_average.update()
+
   def _broadcast_variables(self):
     """Broadcasts variables to other replicas, if required."""
     return
@@ -412,6 +414,10 @@ class MirroredStrategyTrainer(Trainer):
   def _step(self):
     self._strategy.run(super()._step)
 
+  def _update_moving_average(self, moving_average):
+    with self._strategy.scope():
+      super()._update_moving_average(moving_average)
+
 
 def _summarize_gradients(gradients, should_record):
   # Only compute the gradients global norm when the value is actually recorded.
@@ -450,12 +456,11 @@ class MovingAverage(object):
                               "formula and recommended decay values.")
     self._ema = tf.train.ExponentialMovingAverage(decay, num_updates=step)
     self._variables = variables
-    self.update()
 
   @tf.function
   def update(self):
     """Updates the moving average of the variables."""
-    self._ema.apply(var_list=list(map(misc.get_primary_variable, self._variables)))
+    self._ema.apply(self._variables)
 
   @contextlib.contextmanager
   def shadow_variables(self):
@@ -468,7 +473,7 @@ class MovingAverage(object):
     previous_values = []
     for variable in self._variables:
       previous_values.append(variable.value())
-      variable.assign(self._ema.average(misc.get_primary_variable(variable)))
+      variable.assign(self._ema.average(variable))
     yield
     for previous_value, variable in zip(previous_values, self._variables):
       variable.assign(previous_value)
diff --git a/opennmt/utils/misc.py b/opennmt/utils/misc.py
index 98671943..e13d01f2 100644
--- a/opennmt/utils/misc.py
+++ b/opennmt/utils/misc.py
@@ -72,13 +72,6 @@ def get_variable_name(variable, root, model_key="model"):
       return name
   return None
 
-def get_primary_variable(variable):
-  """If :obj:`variable` is distributed, returns the primary component."""
-  # TODO: use a public API to get the primary variable.
-  if isinstance(variable, tf.distribute.DistributedValues):
-    return variable._primary  # pylint: disable=protected-access
-  return variable
-
 def print_as_bytes(text, stream=None):
   """Prints a string as bytes to non rely on :obj:`stream` default encoding.
 


Test patch:
diff --git a/opennmt/tests/training_test.py b/opennmt/tests/training_test.py
index f358e5f6..f625f1ab 100644
--- a/opennmt/tests/training_test.py
+++ b/opennmt/tests/training_test.py
@@ -27,6 +27,7 @@ class TrainingTest(tf.test.TestCase):
     step = tf.Variable(0, trainable=False, dtype=tf.int64)
     variables = [tf.Variable(1.0), tf.Variable(2.0)]
     moving_average = training.MovingAverage(variables, step)
+    moving_average.update()
     variables[0].assign(3.0)
     variables[1].assign(4.0)
     moving_average.update()
@@ -45,9 +46,15 @@ class TrainingTest(tf.test.TestCase):
       step = tf.Variable(0, trainable=False, dtype=tf.int64)
 
     moving_average = training.MovingAverage(variables, step)
+    with strategy.scope():
+      moving_average.update()
+
     variables[0].assign(3.0)
     variables[1].assign(4.0)
-    moving_average.update()
+
+    with strategy.scope():
+      moving_average.update()
+
     self.assertAllEqual(self.evaluate(variables), [3.0, 4.0])
     with moving_average.shadow_variables():
       self.assertAllClose(self.evaluate(variables), [2.8, 3.8])