tensorflow/python/training/moving_averages.py

# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Maintain moving averages of parameters."""
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import reduce_util as ds_reduce_util
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import cond
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variable_v1
from tensorflow.python.ops import variables
from tensorflow.python.training import slot_creator
from tensorflow.python.util.tf_export import tf_export
from tensorflow.tools.docs import doc_controls


@tf_export("__internal__.train.assign_moving_average", v1=[])
def assign_moving_average(variable, value, decay, zero_debias=True, name=None):
  """Compute the moving average of a variable.

  The moving average of 'variable' updated with 'value' is:
    variable * decay + value * (1 - decay)

  The returned Operation sets 'variable' to the newly computed moving average,
  by performing this subtraction:
     variable -= (1 - decay) * (variable - value)

  Since variables that are initialized to a `0` value will be `0` biased,
  `zero_debias` optionally enables scaling by the mathematically correct
  debiasing factor of
    1 - decay ** num_updates
  See Section 3 of (Kingma et al., 2015) for more details.

  The names of the debias shadow variables, by default, include both the scope
  they were created in and the scope of the variables they debias. They are also
  given a uniquifying-suffix.

  E.g.:

  ```
    with tf.compat.v1.variable_scope('scope1'):
      with tf.compat.v1.variable_scope('scope2'):
        var = tf.compat.v1.get_variable('foo')
        update_1 = tf.assign_moving_average(var, 0.0, 1.0)
        update_2 = tf.assign_moving_average(var, 0.0, 0.9)

    # var.name: 'scope1/scope2/foo'
    # shadow var names: 'scope1/scope2/scope1/scope2/foo/biased'
    #                   'scope1/scope2/scope1/scope2/foo/biased_1'
  ```

  Args:
    variable: A Variable.
    value: A tensor with the same shape as 'variable'.
    decay: A float `Tensor` or float value. The moving average decay.
    zero_debias: A python bool. If true, assume the variable is 0-initialized
      and unbias it, as in (Kingma et al., 2015). See docstring in
        `_zero_debias` for more details.
    name: Optional name of the returned operation.

  Returns:
    A tensor which if evaluated will compute and return the new moving average.

  References:
    Adam - A Method for Stochastic Optimization:
      [Kingma et al., 2015](https://arxiv.org/abs/1412.6980)
      ([pdf](https://arxiv.org/pdf/1412.6980.pdf))
  """
  with ops.name_scope(name, "AssignMovingAvg",
                      [variable, value, decay]) as scope:
    decay = ops.convert_to_tensor(1.0 - decay, name="decay")
    if decay.dtype != variable.dtype.base_dtype:
      decay = math_ops.cast(decay, variable.dtype.base_dtype)

    def update_fn(v, value):
      return state_ops.assign_sub(v, (v - value) * decay, name=scope)

    def update(strategy, v, value):
      if zero_debias:
        return _zero_debias(strategy, v, value, decay)
      else:
        return _update(strategy, v, update_fn, args=(value,))

    replica_context = distribute_lib.get_replica_context()
    if replica_context:
      # In a replica context, we update variable using the mean of value across
      # replicas.
      def merge_fn(strategy, v, value):
        value = strategy.extended.reduce_to(ds_reduce_util.ReduceOp.MEAN, value,
                                            v)
        return update(strategy, v, value)

      return replica_context.merge_call(merge_fn, args=(variable, value))
    else:
      strategy = distribute_lib.get_cross_replica_context()
      return update(strategy, variable, value)


def weighted_moving_average(value,
                            decay,
                            weight,
                            truediv=True,
                            collections=None,
                            name=None):
  """Compute the weighted moving average of `value`.

  Conceptually, the weighted moving average is:
    `moving_average(value * weight) / moving_average(weight)`,
  where a moving average updates by the rule
    `new_value = decay * old_value + (1 - decay) * update`
  Internally, this Op keeps moving average variables of both `value * weight`
  and `weight`.

  Args:
    value: A numeric `Tensor`.
    decay: A float `Tensor` or float value. The moving average decay.
    weight:  `Tensor` that keeps the current value of a weight. Shape should be
      able to multiply `value`.
    truediv:  Boolean, if `True`, dividing by `moving_average(weight)` is
      floating point division.  If `False`, use division implied by dtypes.
    collections:  List of graph collections keys to add the internal variables
      `value * weight` and `weight` to. Defaults to
      `[GraphKeys.GLOBAL_VARIABLES]`.
    name: Optional name of the returned operation. Defaults to
      "WeightedMovingAvg".

  Returns:
    An Operation that updates and returns the weighted moving average.
  """
  # Unlike assign_moving_average, the weighted moving average doesn't modify
  # user-visible variables. It is the ratio of two internal variables, which are
  # moving averages of the updates.  Thus, the signature of this function is
  # quite different than assign_moving_average.
  if collections is None:
    collections = [ops.GraphKeys.GLOBAL_VARIABLES]
  with variable_scope.variable_scope(name, "WeightedMovingAvg",
                                     [value, weight, decay]) as scope:
    value_x_weight_var = variable_scope.get_variable(
        "value_x_weight",
        shape=value.get_shape(),
        dtype=value.dtype,
        initializer=init_ops.zeros_initializer(),
        trainable=False,
        collections=collections)
    weight_var = variable_scope.get_variable(
        "weight",
        shape=weight.get_shape(),
        dtype=weight.dtype,
        initializer=init_ops.zeros_initializer(),
        trainable=False,
        collections=collections)
    numerator = assign_moving_average(
        value_x_weight_var, value * weight, decay, zero_debias=False)
    denominator = assign_moving_average(
        weight_var, weight, decay, zero_debias=False)

    if truediv:
      return math_ops.truediv(numerator, denominator, name=scope.name)
    else:
      return math_ops.divide(numerator, denominator, name=scope.name)


def _update(strategy, var, update_fn, args):
  """Applies updates depending on the context."""
  assert distribute_lib.in_cross_replica_context(), (
      "_update can only be called in cross-replica context")
  if distribute_lib.get_update_replica_id() is not None:
    # Call update_fn on var to delegate the implementation. We expect `var` will
    # do the right thing in update context, e.g, if `var` is a MirroredVariable,
    # it should pick its component variable based on `update_replica_id` and
    # only update that.
    return update_fn(var, *args)
  else:
    return strategy.extended.update(var, update_fn, args)


def _zero_debias(strategy, unbiased_var, value, decay):
  """Compute the delta required for a debiased Variable.

  All exponential moving averages initialized with Tensors are initialized to 0,
  and therefore are biased to 0. Variables initialized to 0 and used as EMAs are
  similarly biased. This function creates the debias updated amount according to
  a scale factor, as in (Kingma et al., 2015).

  To demonstrate the bias the results from 0-initialization, take an EMA that
  was initialized to `0` with decay `b`. After `t` timesteps of seeing the
  constant `c`, the variable have the following value:

  ```
    EMA = 0*b^(t) + c*(1 - b)*b^(t-1) + c*(1 - b)*b^(t-2) + ...
        = c*(1 - b^t)
  ```

  To have the true value `c`, we would divide by the scale factor `1 - b^t`.

  In order to perform debiasing, we use two shadow variables. One keeps track of
  the biased estimate, and the other keeps track of the number of updates that
  have occurred.

  Args:
    strategy: `Strategy` used to create and update variables.
    unbiased_var: A Variable representing the current value of the unbiased EMA.
    value: A Tensor representing the most recent value.
    decay: A Tensor representing `1-decay` for the EMA.

  Returns:
    The amount that the unbiased variable should be updated. Computing this
    tensor will also update the shadow variables appropriately.

  References:
    Adam - A Method for Stochastic Optimization:
      [Kingma et al., 2015](https://arxiv.org/abs/1412.6980)
      ([pdf](https://arxiv.org/pdf/1412.6980.pdf))

  """
  with variable_scope.variable_scope(
      unbiased_var.name[:-len(":0")], values=[unbiased_var, value, decay]):
    with ops.init_scope():
      biased_initializer = init_ops.zeros_initializer()
      local_step_initializer = init_ops.zeros_initializer()

    def _maybe_get_unique(name):
      """Get name for a unique variable, if not `reuse=True`."""
      if variable_scope.get_variable_scope().reuse:
        return name
      vs_vars = [
          x.op.name
          for x in variable_scope.get_variable_scope().global_variables()
      ]
      full_name = variable_scope.get_variable_scope().name + "/" + name
      if full_name not in vs_vars:
        return name
      idx = 1
      while full_name + ("_%d" % idx) in vs_vars:
        idx += 1
      return name + ("_%d" % idx)

    with strategy.extended.colocate_vars_with(unbiased_var):
      biased_var = variable_scope.get_variable(
          _maybe_get_unique("biased"),
          initializer=biased_initializer,
          shape=unbiased_var.get_shape(),
          dtype=unbiased_var.dtype,
          trainable=False)
      local_step = variable_scope.get_variable(
          _maybe_get_unique("local_step"),
          shape=[],
          dtype=unbiased_var.dtype,
          initializer=local_step_initializer,
          trainable=False)

  def update_fn(v, value, biased_var, local_step):
    update_biased = state_ops.assign_sub(biased_var,
                                         (biased_var - value) * decay)
    update_local_step = local_step.assign_add(1)

    # This function gets `1 - decay`, so use `1.0 - decay` in the exponent.
    bias_factor = 1 - math_ops.pow(1.0 - decay, update_local_step)
    return state_ops.assign(
        v, update_biased / bias_factor, name=ops.get_name_scope() + "/")

  return _update(
      strategy, unbiased_var, update_fn, args=(value, biased_var, local_step))


@tf_export("train.ExponentialMovingAverage")
class ExponentialMovingAverage:
  """Maintains moving averages of variables by employing an exponential decay.

  When training a model, it is often beneficial to maintain moving averages of
  the trained parameters.  Evaluations that use averaged parameters sometimes
  produce significantly better results than the final trained values.

  The `apply()` method adds shadow copies of trained variables the first time
  it is called, and maintains a moving average of the trained variables in
  their shadow copies at every additional invocation.
  It should generally be called immediately after creating the model weights,
  and then after each training step.

  The `average()` method gives access to the shadow variables.
  It allows you to use the moving averages in place of the last trained values
  for evaluations, by loading the moving averages into your model via
  `var.assign(ema.average(var))`.
  Additionally, although `ExponentialMovingAverage`
  objects are not directly trackable by checkpoints,
  `average()` returns the moving average variables for your model weights,
  which you can then checkpoint. (There is an example
  of this near the bottom of this docstring).
  So, `average()` is useful when
  building an evaluation model, or when restoring a model from a checkpoint
  file.

  The moving averages are computed using exponential decay.  You specify the
  decay value (as a scalar float value, `Tensor`, or `Variable`) when creating
  the `ExponentialMovingAverage` object.  The shadow variables are initialized
  with the same initial values as the trained variables.  When you run `apply`
  to update the moving averages, each shadow variable is updated with the
  formula:

    `shadow_variable -= (1 - decay) * (shadow_variable - variable)`

  This is mathematically equivalent to the classic formula below, but the use
  of an `assign_sub` op (the `"-="` in the formula) allows concurrent lockless
  updates to the variables:

    `shadow_variable = decay * shadow_variable + (1 - decay) * variable`

  Reasonable values for `decay` are close to 1.0, typically in the
  multiple-nines range: 0.999, 0.9999, etc.

  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.

  Example usage when creating a training model:

  ```python
  # Create variables.
  var0 = tf.Variable(...)
  var1 = tf.Variable(...)
  # ... use the variables to build a training model...

  # Create an ExponentialMovingAverage object
  ema = tf.train.ExponentialMovingAverage(decay=0.9999)

  # The first `apply` creates the shadow variables that hold the moving averages
  ema.apply([var0, var1])

  # grab the moving averages for checkpointing purposes or to be able to
  # load the moving averages into the model weights
  averages = [ema.average(var0), ema.average(var1)]

  ...
  def train_step(...):
  ...
    # Apply the optimizer.
    opt.minimize(my_loss, [var0, var1])

    # Update the moving averages
    # of var0 and var1 with additional calls to `apply`
    ema.apply([var0, var1])

  ...train the model by running train_step multiple times...
  ```

  There are several ways to use the moving averages for evaluations:

  1. Assign the values of the shadow variables to your model variables with
     `Variable.assign(...)` before evaluating your
     model. You can use the `average()`
     method to get the shadow variable for a given variable. To continue
     training after using this approach, make sure to record the unaveraged
     weights and restore them before continuing to train. You can see the
     tensorflow-addons' MovingAverage optimizer's `swap_weights` method for
     one example of how to swap variables efficiently in distributed settings:
     https://github.com/tensorflow/addons/blob/v0.13.0/tensorflow_addons/optimizers/moving_average.py#L151
  2. Make sure to checkpoint out your moving average variables in your
     `tf.train.Checkpoint`. At evaluation time, create your shadow variables and
     use `tf.train.Checkpoint` to restore the moving averages into the shadow
     variables. Then, load the moving averages into the actual model weights via
     `var.assign(moving_avg)`.
  3. Checkpoint out your moving average variables in your `tf.train.Checkpoint`.
     For evaluation, restore your model weights directly from the moving
     averages instead of from the non-averaged weights.
     Caution: If you choose this approach, include only the object-graph paths
     to the averaged path in your checkpoint restore.
     If you point both the unaveraged and averaged paths in a checkpoint
     restore to the same variables, it is hard to reason about whether your
     model will restore the averaged or non-averaged variables.

  Example of saving out then restoring the shadow variable values:

  ```python
  # Create variables.
  var0 = tf.Variable(...)
  var1 = tf.Variable(...)
  # ... use the variables to build a training model...

  # Create an ExponentialMovingAverage object, create the shadow variables,
  # and grab the moving averages for checkpointing purposes.
  # (The ExponentialMovingAverage object itself is not checkpointable)
  ema = tf.train.ExponentialMovingAverage(decay=0.9999)
  ema.apply([var0, var1])
  avg_var0 = ema.average(var0)
  avg_var1 = ema.average(var1)

  # Create a Checkpoint that will manage the model weights and the averages,
  checkpoint = tf.train.Checkpoint(model_weights=[var0, var1],
                                   averaged_weights=[avg_var0, avg_var1])
  ... # Do training

  # Save out the checkpoint including the model weights and the moving averages
  checkpoint.save(...)
  ```

  Restore option: restore all averaged & non-averaged weights, then load
  moving averages into the model via `var.assign()`
  ```python
  # Create variables.
  var0 = tf.Variable(...)
  var1 = tf.Variable(...)
  # ... use the variables to build a training model...

  # Create an ExponentialMovingAverage object, create the shadow variables,
  # and grab the moving averages for checkpoint restore purposes.
  # (The ExponentialMovingAverage object itself is not checkpointable)
  ema = tf.train.ExponentialMovingAverage(decay=0.9999)
  ema.apply([var0, var1])
  avg_var0 = ema.average(var0)
  avg_var1 = ema.average(var1)

  # Create a Checkpoint that will manage the model weights and the averages,
  checkpoint = tf.train.Checkpoint(model_weights=[var0, var1],
                                   averaged_weights=[avg_var0, avg_var1])
  checkpoint.restore(...)
  var0.assign(avg_var0)
  var1.assign(avg_var1)
  # var0 and var1 now hold the moving average values
  ```

  Restore option: Directly restore the moving averages into the model weights.
  ```python
  # Create variables.
  var0 = tf.Variable(...)
  var1 = tf.Variable(...)
  # ... use the variables to build a training model...

  # Create a Checkpoint that will manage two objects with trackable state,
  checkpoint = tf.train.Checkpoint(averaged_weights=[var0, var1])
  checkpoint.restore(...)
  # var0 and var1 now hold the moving average values
  ```
  """

  def __init__(self,
               decay,
               num_updates=None,
               zero_debias=False,
               name="ExponentialMovingAverage"):
    """Creates a new ExponentialMovingAverage object.

    The `apply()` method has to be called to create shadow variables.
    Follow-on calls to the `apply()` method will update the moving averages
    in the shadow variables.
    (In TF 1.x graphs `apply()` will return an update op to update
    the moving averages which must be explicitly run).

    The optional `num_updates` parameter allows one to tweak the decay rate
    dynamically. It is typical to pass the count of training steps, usually
    kept in a variable that is incremented at each step, in which case the
    decay rate is lower at the start of training.  This makes moving averages
    move faster.  If passed, the actual decay rate used is:

      `min(decay, (1 + num_updates) / (10 + num_updates))`

    Args:
      decay: A scalar float value, `Tensor`, or `Variable`. The decay parameter.
      num_updates: Optional count of number of updates applied to variables.
      zero_debias: If `True`, zero debias moving-averages that are initialized
        with tensors. (Note: moving averages may not be initialized with
        non-variable tensors when eager execution is enabled).
      name: String. Optional prefix name to use for the name of ops added in
        `apply()`.
    """
    self._decay = decay
    self._num_updates = num_updates
    self._zero_debias = zero_debias
    self._name = name
    self._averages = {}

  @property
  def name(self):
    """The name of this ExponentialMovingAverage object."""
    return self._name

  def apply(self, var_list=None):
    """Maintains moving averages of variables.

    `var_list` must be a list of `Variable` objects.  This method
    creates shadow variables (holding the moving averages)
    for all elements of `var_list`, and
    updates the moving averages using the current `var_list` values. Shadow
    variables for `Variable` objects are initialized to the variable's initial
    value.

    Shadow variables are created with `trainable=False`. To access them you
    can use the EMA object's `average` method. Note that `EMA` objects are
    not trackable by checkpoints, so if you want to checkpoint or restore the
    moving variables you will need to manually grab the shadow
    variables via `average()` and assign them as `tf.Module` properties or
    directly pass them to your `tf.train.Checkpoint`.

    Note that `apply()` can be called multiple times. When eager execution is
    enabled each call to apply will update the variables once, so this needs to
    be called in a loop.

    In legacy TF 1.x graphs, this method returns an op that updates all
    shadow variables from the current value of their associated variables. In
    TF 1.x graphs without automatically control dependencies this op needs to be
    manually run.

    Args:
      var_list: A list of Variable objects. The variables
        must be of types bfloat16, float16, float32, or float64.
        (In legacy TF 1.x graphs these may be tensors, but this is unsupported
        when eager execution is enabled.)

    Returns:
      An Operation that updates the moving averages.

    Raises:
      TypeError: If the arguments are not an allowed type.
    """
    # TODO(touts): op_scope
    if var_list is None:
      var_list = variables.trainable_variables()
    for v in var_list:
      if (isinstance(v, tensor.Tensor)
          and ops.executing_eagerly_outside_functions()):
        raise TypeError(
            "tf.train.ExponentialMovingAverage does not support non-Variable"
            " tensors when eager execution is enabled.")
    zero_debias_true = set()  # set of vars to set `zero_debias=True`
    for var in var_list:
      if var.dtype.base_dtype not in [
          dtypes.bfloat16, dtypes.float16, dtypes.float32, dtypes.float64
      ]:
        raise TypeError("The variables must be half, float, or double: %s" %
                        var.name)

      if var.ref() not in self._averages:
        # For variables: to lower communication bandwidth across devices we keep
        # the moving averages on the same device as the variables. For other
        # tensors, we rely on the existing device allocation mechanism.
        with ops.init_scope():
          if isinstance(var, variables.Variable):
            with ops.device(var.device):
              initialized_value = cond.cond(
                  variable_v1.is_variable_initialized(var), var.read_value,
                  lambda: var.initial_value)  # pylint: disable=cell-var-from-loop
            avg = slot_creator.create_slot(
                var,
                initialized_value,
                self.name,
                colocate_with_primary=True,
                copy_xla_sharding=True)
            # NOTE(mrry): We only add `tf.Variable` objects to the
            # `MOVING_AVERAGE_VARIABLES` collection.
            ops.add_to_collection(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, var)
          else:
            avg = slot_creator.create_zeros_slot(
                var,
                self.name,
                colocate_with_primary=(var.op.type in [
                    "Variable", "VariableV2", "VarHandleOp"
                ]),
                copy_xla_sharding=True)
            if self._zero_debias:
              zero_debias_true.add(avg.ref())
        self._averages[var.ref()] = avg

    with ops.name_scope(self.name) as scope:
      decay = ops.convert_to_tensor(
          self._decay, dtype=dtypes.float32, name="decay")
      if self._num_updates is not None:
        num_updates = math_ops.cast(
            self._num_updates, dtypes.float32, name="num_updates")
        decay = math_ops.minimum(decay,
                                 (1.0 + num_updates) / (10.0 + num_updates))
      updates = []
      for var in var_list:
        avg = self._averages[var.ref()]
        zero_debias = avg.ref() in zero_debias_true
        updates.append(assign_moving_average(avg, var, decay, zero_debias))
      return control_flow_ops.group(*updates, name=scope)

  def average(self, var):
    """Returns the `Variable` holding the average of `var`.

    Args:
      var: A `Variable` object.

    Returns:
      A `Variable` object or `None` if the moving average of `var`
      is not maintained.
    """
    return self._averages.get(var.ref(), None)

  @doc_controls.do_not_generate_docs
  def average_name(self, var):
    """[Meant for TF1] Returns name of `Variable` holding the average for `var`.

    (Designed to work with legacy `tf.compat.v1.train.Saver`, it is sensitive to
    specific variable names and not recommended for TF2)

    The typical scenario for `ExponentialMovingAverage` is to compute moving
    averages of variables during training, and restore the variables from the
    computed moving averages during evaluations.

    To restore variables, you have to know the name of the shadow variables.
    That name and the original variable can then be passed to a `Saver()` object
    to restore the variable from the moving average value with:
      `saver = tf.compat.v1.train.Saver({ema.average_name(var): var})`

    `average_name()` can be called whether or not `apply()` has been called.

    Args:
      var: A `Variable` object.

    Returns:
      A string: The name of the variable that will be used or was used
      by the `ExponentialMovingAverage class` to hold the moving average of
      `var`.
    """
    if var.ref() in self._averages:
      return self._averages[var.ref()].name[:-len(":0")]
    return ops.get_default_graph().unique_name(
        var.name[:-len(":0")] + "/" + self.name, mark_as_used=False)

  @doc_controls.do_not_generate_docs
  def variables_to_restore(self, moving_avg_variables=None):
    """[Designed for TF 1.x] Returns a map of names to `Variables` to restore.

    (Designed to work with legacy `tf.compat.v1.train.Saver`, sensitive to
    specific variable names and not recommended for TF2)

    If a variable has a moving average, use the moving average variable name as
    the restore name; otherwise, use the variable name.

    For example,

    ```python
      variables_to_restore = ema.variables_to_restore()
      saver = tf.compat.v1.train.Saver(variables_to_restore)
    ```

    Below is an example of such mapping:

    ```
      conv/batchnorm/gamma/ExponentialMovingAverage: conv/batchnorm/gamma,
      conv_4/conv2d_params/ExponentialMovingAverage: conv_4/conv2d_params,
      global_step: global_step
    ```

    Args:
      moving_avg_variables: a list of variables that require to use of the
        moving average variable name to be restored. If None, it will default to
        variables.moving_average_variables() + variables.trainable_variables()

    Returns:
      A map from restore_names to variables. The restore_name is either the
      original or the moving average version of the variable name, depending
      on whether the variable name is in the `moving_avg_variables`.
    """
    name_map = {}
    if moving_avg_variables is None:
      # Include trainable variables and variables which have been explicitly
      # added to the moving_average_variables collection.
      moving_avg_variables = variables.trainable_variables()
      moving_avg_variables += variables.moving_average_variables()
    # Remove duplicates
    moving_avg_variables = set(v.ref() for v in moving_avg_variables)
    # Collect all the variables with moving average,
    for v in moving_avg_variables:
      name_map[self.average_name(v.deref())] = v.deref()
    # Make sure we restore variables without moving averages as well.
    moving_avg_variable_names = set(
        v.deref().name for v in moving_avg_variables)
    for v in list(set(variables.global_variables())):
      if v.name not in moving_avg_variable_names and v.op.name not in name_map:
        name_map[v.op.name] = v
    return name_map