| # 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) |
