wehub-resource-sync
661 lines
29 KiB
Python
661 lines
29 KiB
Python
# Copyright 2021 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 the distributed variables library."""
|
|
|
|
import copy
|
|
import os
|
|
|
|
from absl.testing import parameterized
|
|
from tensorflow.python.checkpoint import checkpoint as trackable_utils
|
|
from tensorflow.python.checkpoint import checkpoint_options
|
|
from tensorflow.python.distribute import collective_all_reduce_strategy
|
|
from tensorflow.python.distribute import combinations
|
|
from tensorflow.python.distribute import distribute_lib
|
|
from tensorflow.python.distribute import distribute_utils
|
|
from tensorflow.python.distribute import packed_distributed_variable as packed
|
|
from tensorflow.python.distribute import parameter_server_strategy
|
|
from tensorflow.python.distribute import ps_values
|
|
from tensorflow.python.distribute import strategy_combinations
|
|
from tensorflow.python.distribute import test_util as ds_test_util
|
|
from tensorflow.python.distribute import tpu_strategy
|
|
from tensorflow.python.distribute import values as values_lib
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.eager import test
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import indexed_slices
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import tensor_shape
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import check_ops
|
|
from tensorflow.python.ops import control_flow_assert
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import variable_scope
|
|
from tensorflow.python.ops import variables as variables_lib
|
|
from tensorflow.python.saved_model import save
|
|
from tensorflow.python.saved_model import save_context
|
|
from tensorflow.python.saved_model import save_options
|
|
from tensorflow.python.types import core
|
|
|
|
|
|
def _device_str(d):
|
|
return "/device:GPU:" + str(d)
|
|
|
|
|
|
def _nested_value(d):
|
|
return ("a" + d, ["b" + d, {"c": "d" + d, "e": "f" + d}, "g" + d], "h" + d)
|
|
|
|
|
|
def mirrored_and_tpu_strategy_combinations():
|
|
return combinations.combine(
|
|
distribution=[
|
|
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
|
|
strategy_combinations.mirrored_strategy_with_two_gpus_no_merge_call,
|
|
strategy_combinations.tpu_strategy,
|
|
strategy_combinations.tpu_strategy_packed_var,
|
|
],
|
|
mode=["graph", "eager"])
|
|
|
|
|
|
def checkpoint_test_helper(dvar_test, distribution,
|
|
synchronization, aggregation, enable_async_ckpt):
|
|
# This method is added since `testCheckpointing` cannot be parameterized after
|
|
# the entire class is parameterized.
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
constant_op.constant([1., 2., 3., 4]),
|
|
synchronization=synchronization,
|
|
aggregation=aggregation)
|
|
|
|
dvar_test.evaluate(v.initializer)
|
|
before_save = dvar_test.evaluate(v.read_value())
|
|
|
|
# Save random weights into checkpoint.
|
|
checkpoint = trackable_utils.Checkpoint(v=v)
|
|
ckpt_options = checkpoint_options.CheckpointOptions(
|
|
experimental_enable_async_checkpoint=enable_async_ckpt)
|
|
prefix = os.path.join(dvar_test.get_temp_dir(), "ckpt")
|
|
with dvar_test.test_session():
|
|
save_path = checkpoint.save(file_prefix=prefix, options=ckpt_options)
|
|
|
|
# Assign inverted value.
|
|
dvar_test.evaluate(v.assign(constant_op.constant([4., 3., 2., 1.])))
|
|
after_assign = dvar_test.evaluate(v.read_value())
|
|
dvar_test.assertNotAllClose(before_save, after_assign)
|
|
|
|
# Restore from the checkpoint.
|
|
with dvar_test.test_session():
|
|
checkpoint.restore(save_path).assert_consumed().run_restore_ops()
|
|
after_restore = dvar_test.evaluate(v)
|
|
dvar_test.assertAllClose(before_save, after_restore)
|
|
|
|
# Another round of saving/restoring to ensure that the logic of
|
|
# _copy_trackable_to_cpu works when a copy is already created in object_map.
|
|
dvar_test.evaluate(v.assign(constant_op.constant([5., 6., 7., 8.])))
|
|
before_save_1 = dvar_test.evaluate(v.read_value())
|
|
dvar_test.assertNotAllClose(before_save_1, after_restore)
|
|
with dvar_test.test_session():
|
|
save_path = checkpoint.save(file_prefix=prefix, options=ckpt_options)
|
|
|
|
dvar_test.evaluate(v.assign(constant_op.constant([8., 7., 6., 5.])))
|
|
after_assign_1 = dvar_test.evaluate(v.read_value())
|
|
dvar_test.assertNotAllClose(before_save_1, after_assign_1)
|
|
|
|
with dvar_test.test_session():
|
|
checkpoint.restore(save_path).assert_consumed().run_restore_ops()
|
|
after_restore_1 = dvar_test.evaluate(v)
|
|
dvar_test.assertAllClose(before_save_1, after_restore_1)
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
distribution=[
|
|
strategy_combinations.mirrored_strategy_with_one_cpu,
|
|
strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
|
|
strategy_combinations.mirrored_strategy_with_two_gpus_no_merge_call,
|
|
strategy_combinations.tpu_strategy,
|
|
strategy_combinations.tpu_strategy_packed_var,
|
|
strategy_combinations.tpu_strategy_spmd,
|
|
strategy_combinations.central_storage_strategy_with_gpu_and_cpu,
|
|
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,
|
|
],
|
|
synchronization=[
|
|
variables_lib.VariableSynchronization.ON_READ,
|
|
variables_lib.VariableSynchronization.ON_WRITE,
|
|
],
|
|
aggregation=[
|
|
variables_lib.VariableAggregation.MEAN,
|
|
variables_lib.VariableAggregation.SUM,
|
|
variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
|
|
],
|
|
mode=["graph", "eager"],
|
|
use_var_policy=[True, False]))
|
|
class DistributedVariableTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def testExtendsVariable(self, distribution, synchronization, aggregation):
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
1., synchronization=synchronization, aggregation=aggregation)
|
|
self.assertIsInstance(v, variables_lib.Variable)
|
|
|
|
def testCheckpointing(self, distribution, synchronization, aggregation, mode):
|
|
if (isinstance(distribution,
|
|
collective_all_reduce_strategy.CollectiveAllReduceStrategy)
|
|
and mode == "graph"):
|
|
self.skipTest("MWMS combinations tests do not work well in graph mode.")
|
|
|
|
checkpoint_test_helper(self, distribution, synchronization, aggregation,
|
|
enable_async_ckpt=False)
|
|
|
|
def testAsyncCheckpointing(self, distribution, synchronization,
|
|
aggregation, mode):
|
|
if (isinstance(distribution,
|
|
collective_all_reduce_strategy.CollectiveAllReduceStrategy)
|
|
and mode == "graph"):
|
|
self.skipTest("MWMS combinations tests do not work well in graph mode.")
|
|
|
|
checkpoint_test_helper(self, distribution, synchronization, aggregation,
|
|
enable_async_ckpt=True)
|
|
|
|
def testTraceback(self, distribution, synchronization, aggregation):
|
|
if context.executing_eagerly():
|
|
self.skipTest("does not apply to eager")
|
|
with distribution.scope():
|
|
variable_scope.get_variable(
|
|
name="testVar",
|
|
initializer=1.,
|
|
use_resource=True,
|
|
synchronization=synchronization,
|
|
aggregation=aggregation)
|
|
with self.assertRaisesRegex(ValueError,
|
|
"Variable testVar already exists"):
|
|
variable_scope.get_variable(
|
|
name="testVar",
|
|
initializer=1.,
|
|
use_resource=True,
|
|
synchronization=synchronization,
|
|
aggregation=aggregation)
|
|
|
|
def testSelectReplica(self, distribution, synchronization, aggregation):
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
1., synchronization=synchronization, aggregation=aggregation)
|
|
self.assertIs(v, distribute_utils.select_replica(0, v))
|
|
|
|
def testIsTensorLike(self, distribution, synchronization, aggregation):
|
|
if isinstance(distribution.extended,
|
|
tpu_strategy.TPUExtended) and context.executing_eagerly():
|
|
self.skipTest("TPU doesn't support pure eager")
|
|
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
0., synchronization=synchronization, aggregation=aggregation)
|
|
# In cross replica context.
|
|
self.assertIsInstance(v, core.Tensor)
|
|
# In replica context.
|
|
distribution.run(lambda v: self.assertIsInstance(v, core.Tensor), args=(v,))
|
|
|
|
def testAssignReturnValueIsTensorLike(self, distribution, synchronization,
|
|
aggregation):
|
|
if isinstance(distribution.extended, tpu_strategy.TPUExtended):
|
|
if context.executing_eagerly():
|
|
self.skipTest("TPU doesn't support pure eager")
|
|
else:
|
|
self.skipTest("b/152076846")
|
|
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
0., synchronization=synchronization, aggregation=aggregation)
|
|
|
|
def assert_is_tensor_like(v):
|
|
# We can't use Python literals because they are treated as non-distributed
|
|
# values is not allowed when aggregation is SUM. See
|
|
# `cross_device_ops.reduce_non_distributed_value`.
|
|
delta = array_ops.identity(1.)
|
|
self.assertIsInstance(v.assign(delta), core.Tensor)
|
|
self.assertIsInstance(v.assign_sub(delta), core.Tensor)
|
|
self.assertIsInstance(v.assign_add(delta), core.Tensor)
|
|
|
|
# In cross replica context we return a PerReplica which is not Tensor like
|
|
# all the time yet.
|
|
if (synchronization == variables_lib.VariableSynchronization.ON_READ and
|
|
aggregation != variables_lib.VariableAggregation.SUM):
|
|
assert_is_tensor_like(v)
|
|
|
|
# In replica context.
|
|
distribution.run(assert_is_tensor_like, args=(v,))
|
|
|
|
def testDeepCopy(self, distribution, synchronization, aggregation):
|
|
if not context.executing_eagerly():
|
|
self.skipTest("deepcopy only supported in eager mode")
|
|
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
0., synchronization=synchronization, aggregation=aggregation)
|
|
in_dist_copy = copy.deepcopy(v)
|
|
|
|
out_dist_copy = copy.deepcopy(v)
|
|
|
|
def assert_is_deep_copy(v1, v2):
|
|
self.assertIsInstance(v2, type(v1))
|
|
self.assertEqual(v1.aggregation, v2.aggregation)
|
|
self.assertEqual(v1.distribute_strategy, v2.distribute_strategy)
|
|
if isinstance(v1, ps_values.AggregatingVariable):
|
|
self.assertIsInstance(v2.get(), type(v1.get()))
|
|
self.assertNotEqual(id(v1.get()), id(v2.get()))
|
|
else:
|
|
if v1._policy:
|
|
self.assertNotEqual(id(v1._policy), id(v2._policy)) # pylint: disable=protected-access
|
|
else:
|
|
self.assertEqual(id(v1._policy), id(v2._policy)) # pylint: disable=protected-access
|
|
self.assertEqual(len(v1.values), len(v2.values))
|
|
for (v1v, v2v) in zip(v1.values, v2.values):
|
|
self.assertEqual(v1v.device, v2v.device)
|
|
self.assertNotEqual(id(v1v), id(v2v))
|
|
self.assertAllEqual(
|
|
self.evaluate(v1.values), self.evaluate(v2.values))
|
|
|
|
self.evaluate(variables_lib.global_variables_initializer())
|
|
if not isinstance(distribution.extended, tpu_strategy.TPUExtended):
|
|
distribution.run(assert_is_deep_copy, args=(v, in_dist_copy))
|
|
distribution.run(assert_is_deep_copy, args=(v, out_dist_copy))
|
|
|
|
def testAssignSignature(self, distribution, synchronization, aggregation):
|
|
# This test verifies assign*() can be called in the same way as normal
|
|
# variables.
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
0., synchronization=synchronization, aggregation=aggregation)
|
|
|
|
def assign():
|
|
one = constant_op.constant(1.)
|
|
v.assign(one, True, "assign", False)
|
|
# TODO(b/154017756): SyncOnReadVariable.assign() doesn't support passing
|
|
# value as a keyword argument.
|
|
v.assign(one, use_locking=True, name="assign", read_value=False)
|
|
v.assign_add(one, True, "assign", False)
|
|
v.assign_add(one, use_locking=True, name="assign", read_value=False)
|
|
v.assign_sub(one, True, "assign", False)
|
|
v.assign_sub(one, use_locking=True, name="assign", read_value=False)
|
|
# Return something for graph mode to fetch.
|
|
return constant_op.constant(1)
|
|
|
|
self.evaluate(variables_lib.global_variables_initializer())
|
|
if not (synchronization == variables_lib.VariableSynchronization.ON_READ
|
|
and aggregation == variables_lib.VariableAggregation.SUM):
|
|
self.evaluate(distribution.experimental_local_results(assign()))
|
|
if not (isinstance(distribution.extended, tpu_strategy.TPUExtended) and
|
|
context.executing_eagerly()):
|
|
self.evaluate(
|
|
distribution.experimental_local_results(distribution.run(assign)))
|
|
|
|
def testStrategyExtendedUpdate(self, distribution, synchronization,
|
|
aggregation):
|
|
if len(distribution.extended.parameter_devices) != 2:
|
|
self.skipTest("n/a: needs exactly two parameter devices")
|
|
if (synchronization == variables_lib.VariableSynchronization.ON_WRITE and
|
|
aggregation != variables_lib.VariableAggregation.NONE):
|
|
self.skipTest("n/a: doesn't apply to ON_WRITE variable with aggregation")
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(
|
|
0., synchronization=synchronization, aggregation=aggregation)
|
|
value = values_lib.PerReplica([1., 2.])
|
|
|
|
assign_fn = lambda var, value: var.assign(value)
|
|
self.evaluate(distribution.extended.update(v, assign_fn, args=(value,)))
|
|
self.assertAllEqual(self.evaluate(v.values), [1., 2.])
|
|
|
|
assign_add_fn = lambda var, value: var.assign_add(value)
|
|
self.evaluate(distribution.extended.update(v, assign_add_fn, args=(value,)))
|
|
self.assertAllEqual(self.evaluate(v.values), [2., 4.])
|
|
|
|
assign_sub_fn = lambda var, value: var.assign_sub(value)
|
|
self.evaluate(distribution.extended.update(v, assign_sub_fn, args=(value,)))
|
|
self.assertAllEqual(self.evaluate(v.values), [1., 2.])
|
|
|
|
read_assign_fn = lambda var, value: var.assign_add(var.value() + var.
|
|
read_value())
|
|
self.evaluate(
|
|
distribution.extended.update(v, read_assign_fn, args=(value,)))
|
|
self.assertAllEqual(self.evaluate(v.values), [3., 6.])
|
|
|
|
def testSaveNonDistributed(self, distribution, synchronization, aggregation):
|
|
# This test verifies that the DistributedVariable behave like the primary
|
|
# variable when saving a non-distributed version of the model (the default).
|
|
# The test asserts that the function traced under SaveContext has no device
|
|
# annotations and only reference the primary component of the variable. Note
|
|
# that please avoid capturing other eager tensors in this test to make the
|
|
# assertion easy.
|
|
|
|
if isinstance(distribution.extended,
|
|
parameter_server_strategy.ParameterServerStrategyExtended):
|
|
self.skipTest("b/148689177: AggregatingVariable doesn't "
|
|
"conform to Variable interface well")
|
|
|
|
# tf.function requires the return value to be Tensors, which is not always
|
|
# case for properties and methods of Variable, so we simply discard the
|
|
# return values.
|
|
def _discard_return(f):
|
|
f()
|
|
return
|
|
|
|
def _test(f, v):
|
|
# This verifies that the function under SaveContext:
|
|
# - contains no device annotations.
|
|
# - only references the primary component of the variable.
|
|
g = def_function.function(lambda: _discard_return(f))
|
|
options = save_options.SaveOptions(
|
|
experimental_variable_policy=save_options.VariablePolicy.NONE)
|
|
with save_context.save_context(options):
|
|
# The graph should contain no device.
|
|
graph = g.get_concrete_function().graph
|
|
for op in graph.get_operations():
|
|
self.assertEqual(op.device, "", msg=str(op))
|
|
# The function should only capture the primary variable. Note that it
|
|
# may not have captures, e.g. v.aggregation.
|
|
captures = list(graph.captures)
|
|
self.assertLessEqual(len(captures), 1)
|
|
if graph.captures:
|
|
self.assertIs(captures[0][0], v._primary.handle)
|
|
|
|
def _assert(cond):
|
|
return control_flow_assert.Assert(cond, [cond])
|
|
|
|
with distribution.scope():
|
|
# We use four variables for convenience reasons. They have no special
|
|
# meaning.
|
|
# - v is used whenever possible.
|
|
# - w is used for scatter and gather, which require the variable to be
|
|
# non-scalar.
|
|
# - y is used when the dtype needs to be integer. Note that aggregation
|
|
# cannot be MEAN for integers.
|
|
v = variables_lib.Variable(
|
|
0.,
|
|
synchronization=synchronization,
|
|
aggregation=aggregation,
|
|
trainable=True)
|
|
w = variables_lib.Variable([0., 0., 0.],
|
|
synchronization=synchronization,
|
|
aggregation=aggregation,
|
|
trainable=True)
|
|
if aggregation != variables_lib.VariableAggregation.MEAN:
|
|
y = variables_lib.Variable(
|
|
0, synchronization=synchronization, aggregation=aggregation)
|
|
|
|
# pylint: disable=g-long-lambda
|
|
|
|
# tf.Variable properties.
|
|
_test(lambda: self.assertEqual(v.aggregation, aggregation), v)
|
|
_test(lambda: self.assertIs(v.constraint, None), v)
|
|
# TODO(crccw): should we raise an error instead?
|
|
_test(lambda: self.assertEqual(v.device, v._primary.device), v)
|
|
_test(lambda: self.assertEqual(v.dtype, dtypes.float32), v)
|
|
if not context.executing_eagerly():
|
|
_test(lambda: self.assertIs(v.graph, v._primary.graph), v)
|
|
if not context.executing_eagerly():
|
|
_test(lambda: _assert(v.initial_value == 0), v)
|
|
_test(lambda: self.assertIs(v.initializer, v._primary.initializer), v)
|
|
_test(lambda: self.assertEqual(v.name, "Variable:0"), v)
|
|
if not context.executing_eagerly():
|
|
_test(lambda: self.assertIs(v.op, v._primary.op), v)
|
|
_test(lambda: self.assertEqual(v.shape, tensor_shape.TensorShape(())), v)
|
|
_test(lambda: self.assertEqual(v.synchronization, synchronization), v)
|
|
_test(lambda: self.assertEqual(v.trainable, True), v)
|
|
|
|
# tf.Variable methods.
|
|
_test(lambda: check_ops.assert_equal_v2(v.assign(1.), 1.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v.assign_add(1.), 2.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v.assign_sub(1.), 1.), v)
|
|
# TODO(b/148689177): Implement batch_scatter_update.
|
|
# count_up_to() is skipped since it's deprecated.
|
|
# eval() is skipped since it shouldn't called in a tf.function.
|
|
# experimental_ref() is skipped since it's deprecated.
|
|
# from_proto() is skipped since it shouldn't called in a tf.function.
|
|
# TODO(b/148689177): Implement gather_nd.
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(v.get_shape(),
|
|
tensor_shape.TensorShape(())), v)
|
|
# initialized_value() is skipped since it shouldn't called in a tf.function.
|
|
# load() is skipped since it shouldn't called in a tf.function.
|
|
_test(lambda: check_ops.assert_equal_v2(v.read_value(), 1.), v)
|
|
# ref() is skipped since it shouldn't called in a tf.function.
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
w.scatter_add(_make_index_slices(values=[1., 2.], indices=[0, 2])),
|
|
[1., 0., 2.]), w)
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
w.scatter_div(_make_index_slices(values=[4., 2.], indices=[0, 2])),
|
|
[0.25, 0., 1.]), w)
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
w.scatter_max(_make_index_slices(values=[1., 0.5], indices=[1, 2])),
|
|
[0.25, 1., 1.]), w)
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
w.scatter_min(_make_index_slices(values=[1., 0.5], indices=[0, 1])),
|
|
[0.25, 0.5, 1.]), w)
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
w.scatter_mul(_make_index_slices(values=[2., 0.5], indices=[0, 1])),
|
|
[0.5, 0.25, 1.]), w)
|
|
# TODO(b/148689177): Implement scatter_nd_*
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
w.scatter_sub(_make_index_slices(values=[2., 0.5], indices=[0, 1])),
|
|
[-1.5, -0.25, 1.]), w)
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
w.scatter_update(
|
|
_make_index_slices(values=[2., 0.5], indices=[0, 1])),
|
|
[2., 0.5, 1.]), w)
|
|
# set_shape() is skipped since ResourceVariable doesn't implement it.
|
|
# to_proto() is skipped since it shouldn't called in a tf.function.
|
|
_test(lambda: check_ops.assert_equal_v2(v.value(), 1.), v)
|
|
|
|
# DistributedVariable should be treated as ResourceVariable, so it needs to
|
|
# conform to ResourceVariable interface as well.
|
|
_test(lambda: self.assertIs(v.handle, v._primary.handle), v)
|
|
|
|
# Convert to tensor.
|
|
_test(lambda: check_ops.assert_equal_v2(ops.convert_to_tensor(v), 1.), v)
|
|
|
|
# Control dependency.
|
|
def _with_control_dep():
|
|
with ops.control_dependencies([v.assign(1.)]):
|
|
return array_ops.identity(1)
|
|
|
|
_test(_with_control_dep, v)
|
|
|
|
# Operator overloads.
|
|
_test(lambda: check_ops.assert_equal_v2(v.assign(7.), 7.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v + 1., 8.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(3 + v, 10.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v + v, 14.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v - 2., 5.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v - v, 0.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v * 2., 14.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(3 * v, 21.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(v * v, 49.), v)
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
math_ops.cast(v / 2., dtypes.float32), 3.5), v)
|
|
_test(
|
|
lambda: check_ops.assert_equal_v2(
|
|
math_ops.cast(14. / v, dtypes.float32), 2.), v)
|
|
_test(lambda: _assert(v < 12.), v)
|
|
_test(lambda: _assert(v <= 12.), v)
|
|
_test(lambda: _assert(not v > 12.), v)
|
|
_test(lambda: _assert(not v >= 12.), v)
|
|
_test(lambda: _assert(not 12. < v), v)
|
|
_test(lambda: _assert(not 12. <= v), v)
|
|
_test(lambda: _assert(12. > v), v)
|
|
_test(lambda: _assert(12. >= v), v)
|
|
_test(lambda: check_ops.assert_near_v2(pow(v, 3.), 343.), v)
|
|
_test(lambda: check_ops.assert_near_v2(pow(2., v), 128.), v)
|
|
_test(lambda: check_ops.assert_equal_v2(abs(v), 7.), v)
|
|
|
|
# Operator overloads that only works for integers.
|
|
if aggregation != variables_lib.VariableAggregation.MEAN:
|
|
_test(lambda: check_ops.assert_equal_v2(y.assign(7), 7), y)
|
|
_test(lambda: check_ops.assert_equal_v2(y // 2, 3), y)
|
|
_test(lambda: check_ops.assert_equal_v2(15 // y, 2), y)
|
|
_test(lambda: check_ops.assert_equal_v2(y % 2, 1), y)
|
|
_test(lambda: check_ops.assert_equal_v2(16 % y, 2), y)
|
|
_test(lambda: check_ops.assert_equal_v2(y & 3, 3), y)
|
|
_test(lambda: check_ops.assert_equal_v2(3 & y, 3), y)
|
|
_test(lambda: check_ops.assert_equal_v2(y | 8, 15), y)
|
|
_test(lambda: check_ops.assert_equal_v2(16 | y, 23), y)
|
|
_test(lambda: check_ops.assert_equal_v2(y ^ 3, 4), y)
|
|
_test(lambda: check_ops.assert_equal_v2(11 ^ y, 12), y)
|
|
_test(lambda: check_ops.assert_equal_v2(-y, -7), y)
|
|
_test(lambda: check_ops.assert_equal_v2(~y, ~7), y)
|
|
|
|
# Index.
|
|
if isinstance(distribution.extended, tpu_strategy.TPUExtended):
|
|
# TODO(b/161572567): slice assignment doesn't work for TPU.
|
|
_test(lambda: check_ops.assert_equal_v2(w[0], 2.), w)
|
|
else:
|
|
_test(lambda: check_ops.assert_equal_v2(w[0].assign(1.), [1., 0.5, 1.]),
|
|
w)
|
|
_test(lambda: check_ops.assert_equal_v2(w[0], 1.), w)
|
|
|
|
# pylint: enable=g-long-lambda
|
|
|
|
def testUnsaveable(self, distribution, synchronization, aggregation, mode):
|
|
if isinstance(distribution.extended,
|
|
parameter_server_strategy.ParameterServerStrategyExtended):
|
|
self.skipTest("n/a: not applicable to AggregatingVariable")
|
|
if (isinstance(distribution,
|
|
collective_all_reduce_strategy.CollectiveAllReduceStrategy)
|
|
and mode == "graph"):
|
|
self.skipTest("MWMS combinations tests do not work well in graph mode.")
|
|
if not distribution.extended._use_merge_call():
|
|
self.skipTest("Unsupported combination.")
|
|
with distribution.scope():
|
|
v = variables_lib.Variable([1., 1.],
|
|
synchronization=synchronization,
|
|
aggregation=aggregation)
|
|
|
|
with self.cached_session():
|
|
self.evaluate(variables_lib.global_variables_initializer())
|
|
|
|
export_dir = self.get_temp_dir()
|
|
|
|
def _assert_unsaveable(f):
|
|
# Ignore if it cannot be traced. Certain combinations are not supported or
|
|
# yet or not allowed.
|
|
try:
|
|
f = def_function.function(f).get_concrete_function()
|
|
except (NotImplementedError, ValueError):
|
|
return
|
|
with self.assertRaisesRegex(ValueError, "f_with_input_signature"):
|
|
save.save(v, export_dir, signatures=f)
|
|
|
|
_assert_unsaveable(lambda: v.assign(ops.convert_to_tensor([1., 1.])))
|
|
_assert_unsaveable(lambda: v.assign_add(ops.convert_to_tensor([1., 1.])))
|
|
_assert_unsaveable(lambda: v.assign_sub(ops.convert_to_tensor([1., 1.])))
|
|
_assert_unsaveable(lambda: v.scatter_add(_make_index_slices([1.], [0])))
|
|
_assert_unsaveable(lambda: v.scatter_sub(_make_index_slices([1.], [0])))
|
|
_assert_unsaveable(lambda: v.scatter_mul(_make_index_slices([1.], [0])))
|
|
_assert_unsaveable(lambda: v.scatter_div(_make_index_slices([1.], [0])))
|
|
_assert_unsaveable(lambda: v.scatter_min(_make_index_slices([1.], [0])))
|
|
_assert_unsaveable(lambda: v.scatter_max(_make_index_slices([1.], [0])))
|
|
_assert_unsaveable(lambda: v.scatter_update(_make_index_slices([1.], [0])))
|
|
# Reading a ON_READ variable should be unsaveable if either:
|
|
# 1) CollectiveAllReduceStrategy, and aggregation is MEAN/SUM.
|
|
# 2) aggregation is SUM.
|
|
if (synchronization == variables_lib.VariableSynchronization.ON_READ and
|
|
(aggregation == variables_lib.VariableAggregation.SUM or
|
|
(not distribution.extended._use_merge_call()) or
|
|
(isinstance(distribution.extended,
|
|
collective_all_reduce_strategy.CollectiveAllReduceExtended)
|
|
and aggregation == variables_lib.VariableAggregation.MEAN))):
|
|
_assert_unsaveable(v.read_value)
|
|
_assert_unsaveable(v.value)
|
|
_assert_unsaveable(lambda: ops.convert_to_tensor(v))
|
|
else:
|
|
# Otherwise reading a variable should be saveable.
|
|
|
|
@def_function.function
|
|
def f():
|
|
v.read_value()
|
|
v.value()
|
|
return ops.convert_to_tensor(v)
|
|
|
|
with self.cached_session():
|
|
save.save(v, export_dir, signatures=f.get_concrete_function())
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
distribution=[
|
|
strategy_combinations.mirrored_strategy_with_one_cpu,
|
|
strategy_combinations.tpu_strategy,
|
|
],
|
|
mode=["eager"]))
|
|
class PackedDistributedVariableTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def testPackedVariable(self, distribution):
|
|
with distribution.scope():
|
|
v0 = variables_lib.Variable(0.)
|
|
self.assertIsNone(v0._packed_var)
|
|
|
|
distribution._enable_packed_variable_in_eager_mode = True
|
|
with distribution.scope():
|
|
v1 = variables_lib.Variable(0)
|
|
self.assertIsInstance(v1._packed_var, packed.PackedDistributedVariable)
|
|
|
|
devices = v1._devices
|
|
for i in range(1, len(devices)):
|
|
with distribute_lib.ReplicaContext(distribution, i):
|
|
v1.assign(i)
|
|
val = v1._get()
|
|
self.assertIsInstance(val, packed.PackedVarAndDevice)
|
|
self.assertEqual(val.device, devices[0])
|
|
self.assertEqual(self.evaluate(val.read_value()), 0)
|
|
for i in range(0, len(devices)):
|
|
with distribute_lib.ReplicaContext(distribution, i):
|
|
val = v1._get()
|
|
self.assertIsInstance(val, packed.PackedVarAndDevice)
|
|
self.assertEqual(val.device, devices[i])
|
|
self.assertEqual(self.evaluate(val.read_value()), i)
|
|
|
|
def testIgnorePackedVariableInSaveContext(self, distribution):
|
|
distribution._enable_packed_variable_in_eager_mode = True
|
|
with distribution.scope():
|
|
v = variables_lib.Variable(0)
|
|
self.assertIsInstance(v._packed_variable,
|
|
packed.PackedDistributedVariable)
|
|
|
|
options = save_options.SaveOptions()
|
|
with save_context.save_context(options):
|
|
self.assertIsNone(v._packed_variable)
|
|
|
|
|
|
def _make_index_slices(values, indices, dense_shape=None):
|
|
if dense_shape:
|
|
dense_shape = array_ops.identity(dense_shape)
|
|
return indexed_slices.IndexedSlices(
|
|
array_ops.identity(values), array_ops.identity(indices), dense_shape)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ds_test_util.main()
|