moving_averages_test.py

# Copyright 2018 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.
# ==============================================================================
"""Tests for training.moving_averages when using a DistributionStrategy."""

from absl.testing import parameterized

from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.distribute import strategy_test_lib
from tensorflow.python.distribute import test_util
from tensorflow.python.eager import def_function
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.ops import variables
from tensorflow.python.training import moving_averages


all_distributions = [
    strategy_combinations.default_strategy,
    strategy_combinations.one_device_strategy,
    strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
    strategy_combinations.central_storage_strategy_with_gpu_and_cpu,
    strategy_combinations.tpu_strategy,
    strategy_combinations.multi_worker_mirrored_2x1_cpu,
    strategy_combinations.multi_worker_mirrored_2x1_gpu,
    strategy_combinations.multi_worker_mirrored_2x2_gpu,
    strategy_combinations.multi_worker_mirrored_2x2_gpu_no_merge_call,
    strategy_combinations.multi_worker_mirrored_4x1_cpu,
]

all_combinations = combinations.combine(
    distribution=all_distributions, mode=["graph"])

all_combinations_eager = combinations.combine(
    distribution=all_distributions, mode=["eager"], use_function=[True, False])


class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase):

  @combinations.generate(all_combinations)
  def testReplicaModeWithoutZeroDebias(self, distribution):
    replica_id = [0]

    def replica_fn():
      var = variables.Variable([10.0, 11.0])
      val = constant_op.constant([1.0 + replica_id[0], 2.0 - replica_id[0]])
      replica_id[0] += 1
      decay = 0.25
      assign = moving_averages.assign_moving_average(
          var, val, decay, zero_debias=False)
      return var, assign

    with distribution.scope():
      var, assign = distribution.extended.call_for_each_replica(replica_fn)
      self.evaluate(variables.global_variables_initializer())
      self.assertAllClose([10.0, 11.0], self.evaluate(var))
      self.evaluate(distribution.experimental_local_results(assign))
      # Mean of val across calls to replica_fn().
      average_val = [1.0 + 0.5 * (replica_id[0] - 1),
                     2.0 - 0.5 * (replica_id[0] - 1)]
      val_weight = 1.0 - 0.25
      self.assertAllClose(
          [10.0 * 0.25 + average_val[0] * val_weight,
           11.0 * 0.25 + average_val[1] * val_weight],
          self.evaluate(var))

  @combinations.generate(all_combinations)
  def testReplicaMode(self, distribution):
    replica_id = [0]

    def replica_fn():
      var = variables.Variable([0.0, 0.0])
      val = constant_op.constant([1.0 + replica_id[0], 2.0 - replica_id[0]])
      replica_id[0] += 1
      decay = 0.25
      assign = moving_averages.assign_moving_average(var, val, decay)
      return var, assign.op

    with distribution.scope():
      var, assign_op = distribution.extended.call_for_each_replica(replica_fn)
      self.evaluate(variables.global_variables_initializer())
      self.assertAllClose([0.0, 0.0], self.evaluate(var))
      self.evaluate(distribution.experimental_local_results(assign_op))
      # Mean of val across calls to replica_fn().
      average_val = [1.0 + 0.5 * (replica_id[0] - 1),
                     2.0 - 0.5 * (replica_id[0] - 1)]
      self.assertAllClose(average_val, self.evaluate(var))

  @combinations.generate(all_combinations)
  def testCrossDeviceWithoutZeroDebias(self, distribution):
    with distribution.scope():
      var = variables.Variable([10.0, 11.0])
      val = constant_op.constant([1.0, 2.0])
      decay = 0.25
      # NOTE(josh11b): We currently generate an error if val is a PerReplica
      # value.
      assign = moving_averages.assign_moving_average(
          var, val, decay, zero_debias=False)

      self.evaluate(variables.global_variables_initializer())
      self.assertAllClose([10.0, 11.0], self.evaluate(var))
      self.evaluate(assign)
      average_val = [1.0, 2.0]
      val_weight = 1.0 - 0.25
      self.assertAllClose(
          [10.0 * 0.25 + average_val[0] * val_weight,
           11.0 * 0.25 + average_val[1] * val_weight],
          self.evaluate(var))
      # Also try assign.op.
      self.evaluate(assign.op)
      orig_weight = 0.25 * 0.25
      val_weight = 1.0 - orig_weight
      self.assertAllClose(
          [10.0 * orig_weight + average_val[0] * val_weight,
           11.0 * orig_weight + average_val[1] * val_weight],
          self.evaluate(var))

  @combinations.generate(all_combinations)
  def testCrossDevice(self, distribution):
    with distribution.scope():
      var = variables.Variable([0.0, 0.0])
      val = variables.Variable([1.0, 2.0])
      decay = 0.25
      # NOTE(josh11b): We currently generate an error if val is a PerReplica
      # value.
      assign = moving_averages.assign_moving_average(var, val, decay)

      self.evaluate(variables.global_variables_initializer())
      self.assertAllClose([0.0, 0.0], self.evaluate(var))
      self.evaluate(assign)
      self.assertAllClose([1.0, 2.0], self.evaluate(var))

  @combinations.generate(all_combinations_eager)
  def testUpdateContext(self, distribution, use_function):
    with distribution.scope():
      var1 = variables.Variable([0.0, 0.0])
      var2 = variables.Variable([0.0, 0.0])
      var3 = variables.Variable([0.0, 0.0])

      def update_fn(v, value):
        v.assign_add(value)
        moving_averages.assign_moving_average(var2, [2.0, 4.0], decay=0.25)
        moving_averages.assign_moving_average(
            var3, [2.0, 4.0], decay=0.25, zero_debias=False)

      distribution.extended.update(var1, update_fn, ([1.0, 1.0],))

      self.assertAllClose([2.0, 4.0], var2.read_value())
      self.assertAllClose([1.5, 3.0], var3.read_value())

  @combinations.generate(all_combinations)
  def testAssignVariable(self, distribution):

    def replica_fn():
      var = variables.Variable([10.0, 11.0])
      # Here we expect to check the case when input value are variable.
      val = variables.Variable([1., 2.])
      decay = 0.25
      assign = moving_averages.assign_moving_average(
          var, val, decay, zero_debias=False)
      return var, assign

    with distribution.scope():
      var, assign = distribution.extended.call_for_each_replica(replica_fn)
      self.evaluate(variables.global_variables_initializer())
      self.assertAllClose([10.0, 11.0], self.evaluate(var))
      self.evaluate(distribution.experimental_local_results(assign))
      self.assertAllClose(
          [10 * 0.25 + 1. * (1 - 0.25), 11 * 0.25 + 2. * (1 - 0.25)],
          self.evaluate(var))


class ExponentialMovingAverageTest(test.TestCase, parameterized.TestCase):

  @combinations.generate(all_combinations_eager)
  def testReplicaContextEager(self, distribution, use_function):
    if not use_function and strategy_test_lib.is_tpu_strategy(distribution):
      self.skipTest("TPUStrategy doesn't support pure eager execution.")
    with distribution.scope():
      w = variables.Variable([1.0],
                             name="w",
                             aggregation=variables.VariableAggregation.MEAN)
      ema = moving_averages.ExponentialMovingAverage(0.8)

      def fn():

        def _ema_replica_fn_eager():
          ema.apply([w])
          w.assign_sub([0.5])
          ema.apply([w])
          return ema.average(w)

        return distribution.run(_ema_replica_fn_eager)

      if use_function:
        fn = def_function.function(fn)
      ema_w = fn()
    self.assertAllClose(
        self.evaluate(distribution.experimental_local_results(ema_w))[0],
        [0.89999998])

  @combinations.generate(all_combinations_eager)
  def testCrossReplicaContextEager(self, distribution, use_function):
    with distribution.scope():
      w = variables.Variable([1.0],
                             name="w",
                             aggregation=variables.VariableAggregation.MEAN)
      ema = moving_averages.ExponentialMovingAverage(0.8)

      def fn():
        ema.apply([w])
        w.assign_sub([0.5])
        ema.apply([w])
        return ema.average(w)

      if use_function:
        fn = def_function.function(fn)
      avg = fn()
    self.assertAllClose(
        self.evaluate(distribution.experimental_local_results(avg))[0],
        [0.89999998])

  def _ema_replica_fn_graph(self):
    w = variables.Variable([1.0],
                           name="w",
                           aggregation=variables.VariableAggregation.MEAN)
    ema = moving_averages.ExponentialMovingAverage(0.8)
    w_apply = ema.apply([w])
    w_assign = w.assign_sub([0.5])
    return w_assign, w_apply, ema.average(w)

  @combinations.generate(all_combinations)
  def testReplicaContextGraph(self, distribution):
    if strategy_test_lib.is_tpu_strategy:
      self.skipTest("b/139550827: Cannot do variable.assign in replica context "
                    "of TPUStrategy")
    with distribution.scope():
      w_assign, w_apply, ema_w = distribution.run(
          self._ema_replica_fn_graph)
    self.assertEqual(ema_w.name, "w/ExponentialMovingAverage:0")
    self.evaluate(variables.global_variables_initializer())
    self.evaluate(distribution.experimental_local_results(w_apply))
    self.evaluate(distribution.experimental_local_results(w_assign))
    self.evaluate(distribution.experimental_local_results(w_apply))
    self.assertAllClose(
        self.evaluate(distribution.experimental_local_results(ema_w))[0],
        [0.89999998])

  @combinations.generate(all_combinations)
  def testCrossReplicaContextGraph(self, distribution):
    with distribution.scope():
      w_assign, w_apply, ema_w = self._ema_replica_fn_graph()
    self.assertEqual(ema_w.name, "w/ExponentialMovingAverage:0")
    self.evaluate(variables.global_variables_initializer())
    self.evaluate(distribution.experimental_local_results(w_apply))
    self.evaluate(distribution.experimental_local_results(w_assign))
    self.evaluate(distribution.experimental_local_results(w_apply))
    self.assertAllClose(
        self.evaluate(distribution.experimental_local_results(ema_w))[0],
        [0.89999998])


if __name__ == "__main__":
  # TODO(b/172304955): enable logical devices.
  test_util.main(config_logical_devices=False)