Implement EMA schedule

tf.train.ExponentialMovingAverage [1] accepts decay=tf.Variable.
To have fine-grained control over the value of the decay parameter during
training, pass a scalar tf.Variable as the decay value to the constructor, and
update the variable as needed.
[1] https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage

For example, we can easily have pure weight average over t, as setting
ema_schedule to t**-1. It give average theta = mean(theta_i, (0, t)).

PiperOrigin-RevId: 490081739

Author: ds-hwang

lingvo/core/base_model.py

@@ -18,7 +18,7 @@
 import dataclasses
 import functools
 import re
-from typing import Dict, Union
+from typing import Dict, Optional, Union
 
 import lingvo.compat as tf
 from lingvo.core import base_input_generator
@@ -64,6 +64,15 @@ class DecodeEmailOptions:
   global_step: int
 
 
+@dataclasses.dataclass(frozen=True)
+class ExecutorEma:
+  """EMA related instances which an executor prepares."""
+  # EMA object.
+  ema: Optional[tf.train.ExponentialMovingAverage] = None
+  # ema_decay variable.
+  ema_decay: Optional[tf.Variable] = None
+
+
 def _VariablesForEMA(params, model_var_list):
   """Gets a list of variables that need to apply exponential moving average."""
   # Use variable reference since variable is not hashable in eager mode.
@@ -168,6 +177,7 @@ def Params(cls):
     tp.Define(
         'ema_decay_moving_vars', None,
         'If True, include variables from collection "moving_vars" in ema.')
+    tp.Define('ema_schedule', None, 'EMA decay schedule over global_step.')
     tp.Define(
         'init_from_checkpoint_rules', {},
         'If not None, a dictionary with keys corresponding to a checkpoint '
@@ -436,6 +446,9 @@ def __init__(self, params):
             self.CreateChildren('learners', tp.learner)
           else:
             self.CreateChildren('learners', [tp.learner])
+
+        if tp.ema_schedule:
+          self.CreateChild('ema_schedule', tp.ema_schedule)
       self._UpdateVnConfig()
 
       if (tp and tp.pruning_hparams_dict and
@@ -857,9 +870,22 @@ def ApplyExponentialMovingAverage(self):
     self._graphs_applied_ema.add(graph)
 
     tf.logging.info('ApplyExponentialMovingAverage on %s', self)
+    pre_op = tf.no_op()
+
+    # Update EMA decay.
+    tp = self.params.train
+    if tp.ema_schedule:
+      assert isinstance(self.parent, BaseModel)
+      ema_decay_var = self.parent.ema_decay
+      assert ema_decay_var is not None
+      ema_decay = self.ema_schedule.Value(step=self._global_step_var)
+      ema_decay = tf.minimum(ema_decay, 1.0)
+      pre_op = ema_decay_var.assign(ema_decay, read_value=False)
+
     # Use empty name here so no prefix is added to the EMA variable names.
     scoped_creator = py_utils.GetLingvoVariableCreator('', '')
-    return scoped_creator(self._ApplyEMA, ema=ema)
+    with tf.control_dependencies([pre_op]):
+      return scoped_creator(self._ApplyEMA, ema=ema)
 
   def _ApplyEMA(self, ema):
     all_vars = _VariablesForEMA(self.params, self.vars.Flatten())
@@ -1134,6 +1160,7 @@ def Params(cls):
         'ema_decay_moving_vars', None,
         'If True, include variables from collection "moving_vars" in ema. '
         'Must be set consistent across all tasks.')
+    tp.Define('ema_schedule', None, 'EMA decay schedule over global_step.')
     tp.Define('init_from_checkpoint_rules', {},
               'See BaseTask documentation for details.')
     tp.Define('init_from_checkpoint_override', '',
@@ -1163,7 +1190,7 @@ def Params(cls):
         'checkpoints. Currently only support custom saver.')
     return p
 
-  def __init__(self, params, executor_ema=None):
+  def __init__(self, params, executor_ema=ExecutorEma()):
     """Initializes this Model."""
     assert issubclass(params.cls, BaseModel)
     super().__init__(params)
@@ -1172,21 +1199,23 @@ def __init__(self, params, executor_ema=None):
     self._global_step_var = py_utils.GetOrCreateGlobalStepVar()
 
     tp = self.params.train
-    if tp.ema_decay > 0:
+    if tp.ema_decay > 0 or tp.ema_schedule:
       assert tp.ema_decay < 1.0
-      assert self.cluster.is_executor_tpu == (executor_ema is not None)
-      if executor_ema is not None:
+      assert self.cluster.is_executor_tpu == (executor_ema.ema is not None)
+      if executor_ema.ema is not None:
         # Use the EMA for executor training if set.
-        self._ema = executor_ema
+        self._ema, self._ema_decay = executor_ema.ema, executor_ema.ema_decay
       else:
-        self._ema = py_utils.CreateEMAForModel(self.params, self.global_step)
+        self._ema_decay = py_utils.CreateEMADecayVar(self.params)
+        self._ema = py_utils.CreateEMAForModel(self.params, self.global_step,
+                                               self._ema_decay)
     else:
       # Evaler/Decoder may disable EMA while ExecutorTpu uses EMA. executor_ema
       # depends on the trainer task params, but Evaler/Decoder may have
       # different task params (e.g. ema_decay=0). See model_registry.py
       if not self.do_eval:
-        assert not executor_ema
-      self._ema = None
+        assert executor_ema.ema is None
+      self._ema = self._ema_decay = None
     self._ema_variables_dict = {}
 
   @property
@@ -1197,6 +1226,10 @@ def global_step(self):
   def variables_for_ema(self):
     return _VariablesForEMA(self.params, self.vars.Flatten())
 
+  @property
+  def ema_decay(self):
+    return self._ema_decay
+
   def _MakeEMAVariablesDict(self):
     if self.ema:
       res = {}
@@ -1364,6 +1397,7 @@ def CopyTaskParams(cls, task_params, p):
     tp.checkpoint_finite_check = p.task.train.checkpoint_finite_check
     tp.ema_decay = p.task.train.ema_decay
     tp.ema_decay_moving_vars = p.task.train.ema_decay_moving_vars
+    tp.ema_schedule = p.task.train.ema_schedule
 
   def __init__(self, params, **kwargs):
     assert issubclass(params.cls, SingleTaskModel)
@@ -1498,10 +1532,11 @@ def __init__(self, params, **kwargs):
       p.task_schedule = task_scheduler.ConstantScheduler.Params()
       p.task_schedule.task_probs = sorted(list(p.task_probs.IterParams()))
 
-    if p.train.ema_decay > 0:
+    tp = p.train
+    if tp.ema_decay > 0 or tp.ema_schedule:
       for task_name, task_params in sorted_task_params:
         for field in ['ema_decay', 'ema_decay_moving_vars']:
-          if task_params.train.Get(field) != p.train.Get(field):
+          if task_params.train.Get(field) != tp.Get(field):
             raise ValueError('Params did not match for field %s in task %s' %
                              (field, task_name))
 

lingvo/core/ema_test.py

@@ -23,6 +23,7 @@
 from lingvo.core import cluster_factory
 from lingvo.core import layers
 from lingvo.core import py_utils
+from lingvo.core import schedule
 from lingvo.core import test_utils
 import mock
 import numpy as np
@@ -146,8 +147,11 @@ def testBatchNormLayer(self):
         graph=tf.Graph()) as sess, cluster_factory.ForTestingWorker(
             job='executor_tpu', do_eval=True), mock.patch(
                 'lingvo.core.py_utils.use_tpu', return_value=True):
-      executor_ema = py_utils.CreateEMAForModel(
-          p, py_utils.GetOrCreateGlobalStepVar())
+      ema_decay_var = None
+      ema_var = py_utils.CreateEMAForModel(p,
+                                           py_utils.GetOrCreateGlobalStepVar(),
+                                           ema_decay_var)
+      executor_ema = base_model.ExecutorEma(ema_var, ema_decay_var)
       model = p.Instantiate(executor_ema=executor_ema)
       self.assertIsNotNone(model.ema)
       task = model._task
@@ -166,5 +170,81 @@ def testBatchNormLayer(self):
       self.assertAllClose([beta_1, beta_1_ema, mean_1, mean_1_ema],
                           self.evaluate([beta, beta_ema, mean, mean_ema]))
 
+  def testEmaSchedule(self):
+    task = self.TestParams(layers.BatchNormLayer.Params().Set(dim=1))
+    task.train.ema_decay = 0
+    # Note: EMA = decay * EMA + (1 - decay) * var
+    ema_off = 1.0  # ema keeps constant.
+    ema_is_var = 0.0  # ema copys var value.
+    task.train.ema_schedule = schedule.PiecewiseConstantSchedule.Params().Set(
+        boundaries=[99, 199], values=[ema_off, 0.9, ema_is_var])
+    task.train.ema_decay_moving_vars = True
+    p = base_model.SingleTaskModel.Params(task)
+    model = p.Instantiate()
+    self.assertIsNotNone(model.ema)
+    self.assertIsNotNone(model.ema_decay)
+    task = model._task
+
+    layer = task.encoder
+    self.assertLen(layer.vars, 4)
+    for var in layer.vars.Flatten():
+      self.assertIsNotNone(model.ema.average(var), msg=var.name)
+    beta = layer.vars.beta
+    mean = layer.vars.moving_mean
+
+    beta_0 = np.asarray([0.])
+    mean_0 = np.asarray([0.])
+    beta_1 = np.asarray([.2])
+    mean_1 = np.asarray([.03])
+    beta_1_ema = beta_1 * .1
+    mean_1_ema = mean_1 * .1
+    # Check EMA decay schedul in Train.
+    with self.session():
+      # Test EMA values.
+      self.evaluate(tf.global_variables_initializer())
+      # var is initialized as 0, and EMA assigns the var value.
+      self.assertAllClose([beta_0, beta_0, mean_0, mean_0],
+                          self.evaluate([
+                              beta,
+                              model.ema.average(beta), mean,
+                              model.ema.average(mean)
+                          ]))
+
+      # At step=1, ema_decay=1.0 by ema_schedule. EMA update is off.
+      global_step = 1
+      self.evaluate(tf.assign(py_utils.GetOrCreateGlobalStepVar(), global_step))
+      self.evaluate(tf.assign(beta, beta_1))
+      self.evaluate(tf.assign(mean, mean_1))
+      ema_op = task.ApplyExponentialMovingAverage()
+      self.evaluate(ema_op)
+      self.assertAllClose([beta_1, beta_0, mean_1, mean_0],
+                          self.evaluate([
+                              beta,
+                              model.ema.average(beta), mean,
+                              model.ema.average(mean)
+                          ]))
+
+      # At step=100, ema_decay=0.9 by ema_schedule.
+      global_step = 100
+      self.evaluate(tf.assign(py_utils.GetOrCreateGlobalStepVar(), global_step))
+      self.evaluate(ema_op)
+      self.assertAllClose([beta_1, beta_1_ema, mean_1, mean_1_ema],
+                          self.evaluate([
+                              beta,
+                              model.ema.average(beta), mean,
+                              model.ema.average(mean)
+                          ]))
+
+      # At step=200, ema_decay=0.0 by ema_schedule. EMA copies var value.
+      global_step = 200
+      self.evaluate(tf.assign(py_utils.GetOrCreateGlobalStepVar(), global_step))
+      self.evaluate(ema_op)
+      self.assertAllClose([beta_1, beta_1, mean_1, mean_1],
+                          self.evaluate([
+                              beta,
+                              model.ema.average(beta), mean,
+                              model.ema.average(mean)
+                          ]))
+
 if __name__ == '__main__':
   test_utils.main()

lingvo/core/program.py

@@ -108,7 +108,7 @@ def __init__(self,
                params,
                shared_model=None,
                trial=base_trial.NoOpTrial(),
-               ema=None,
+               executor_ema=base_model.ExecutorEma(),
                **kwargs):
     self.params = params.Copy()
     p = self.params
@@ -121,7 +121,7 @@ def __init__(self,
     self._tf_master = kwargs.pop('tf_master', None)
     self._write_train_input_stats = p.write_train_input_stats
     self._trial = trial
-    self._ema = ema
+    self._executor_ema = executor_ema
 
     self._SetProgramDir()
     # Initialized on use; access via self._summary_writer property only.
@@ -394,8 +394,8 @@ def _InstantiateTaskModel(self, task_params):
     """
     if issubclass(task_params.cls, base_model.MultiTaskSubModel):
       return task_params.Instantiate(
-          shared_model=self._shared_model, executor_ema=self._ema)
-    return task_params.Instantiate(executor_ema=self._ema)
+          shared_model=self._shared_model, executor_ema=self._executor_ema)
+    return task_params.Instantiate(executor_ema=self._executor_ema)
 
   def _OutfeedEnqueue(self, per_example_tensors):
     if not per_example_tensors:
@@ -1875,7 +1875,7 @@ def BuildTpuSubgraph(self):
     self._eval_metrics = metrics.TpuEvalMetrics(max_metrics=p.max_metrics)
     with py_utils.OpportunisticVariableReuseScope(True):
       self._train_model = self._train_task_params.Instantiate(
-          executor_ema=self._ema)
+          executor_ema=self._executor_ema)
     self._train_task = self._train_model.GetTask()
     self._train_task.input.InfeedSetupGraph()
     self._model = self._train_model

lingvo/core/py_utils.py

@@ -751,13 +751,27 @@ def GradientTape(*args, **kwargs):
     _GRADIENT_TAPE_STACK.stack.pop()
 
 
-def CreateEMAForModel(model_params, global_step):
+def CreateEMADecayVar(model_params):
+  """Creates an EMA decay variable."""
+  p = model_params
+  tp = p.train
+  if tp.ema_schedule:
+    tf.logging.log_if(tf.logging.WARNING,
+                      f'ema_schedule overrides ema_decay:{tp.ema_decay}.',
+                      tp.ema_decay > 0)
+    wp = WeightParams(
+        shape=[], init=WeightInit.Constant(tp.ema_decay), dtype=p.dtype)
+    return CreateVariable('ema_decay', wp, trainable=False)
+  return None
+
+
+def CreateEMAForModel(model_params, global_step, ema_decay):
   """Creates an EMA object for model with param `model_params` if applicable."""
   p = model_params
 
   # Check that EMA settings for the model and subtasks match.
   def CheckEMA(task_name, task_params):
-    for field in ['ema_decay', 'ema_decay_moving_vars']:
+    for field in ['ema_decay', 'ema_decay_moving_vars', 'ema_schedule']:
       model_value = p.train.Get(field)
       task_value = task_params.train.Get(field)
       if task_value != model_value:
@@ -774,9 +788,16 @@ def CheckEMA(task_name, task_params):
     # SingleTaskModel.
     CheckEMA(p.task.name, p.task)
 
-  if p.train.ema_decay > 0:
+  tp = p.train
+  if tp.ema_decay > 0 or tp.ema_schedule:
+    if tp.ema_schedule:
+      assert isinstance(ema_decay, tf.Variable)
+      # ema_decay takes all control. Otherwise, global_step affects ema_decay.
+      global_step = None
+    else:
+      ema_decay = p.train.ema_decay
     return tf.train.ExponentialMovingAverage(
-        decay=p.train.ema_decay, num_updates=global_step)
+        decay=ema_decay, num_updates=global_step)
   return None
 
 

lingvo/executor.py

@@ -292,9 +292,13 @@ def _WaitTillInit(job=None):
 
     self._checkpoint_to_load = None
     with self._cluster:
-      # Create the ExponentialMovingAverage singleton shared by all programs, if
-      # applicable.
-      ema = py_utils.CreateEMAForModel(train_cfg, self._global_step_var)
+      with tf.container(self._container_id), contextlib.ExitStack() as stack:
+        if not py_utils.IsEagerMode():
+          stack.enter_context(self._graph.as_default())
+        ema_decay_var = py_utils.CreateEMADecayVar(train_cfg)
+      ema_obj = py_utils.CreateEMAForModel(train_cfg, self._global_step_var,
+                                           ema_decay_var)
+      executor_ema = base_model.ExecutorEma(ema_obj, ema_decay_var)
       tf.logging.info('ps_params_dict=%s',
                       {k: v.ToText() for k, v in ps_params_dict.items()})
       for task_string, program_schedule_params in ps_params_dict.items():
@@ -306,7 +310,7 @@ def _WaitTillInit(job=None):
         ps = program_schedule_params.Instantiate(
             shared_model=shared_model,
             trial=self._trial,
-            ema=ema,
+            executor_ema=executor_ema,
             tf_master=self._tf_master)
         self._program_schedule_dict[task_string] = ps
         tf.logging.info('program_schedule_params: %s',