wehub-resource-sync
481 lines
16 KiB
Python
481 lines
16 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 values library."""
|
|
|
|
from absl.testing import parameterized
|
|
|
|
from tensorflow.python.distribute import combinations
|
|
from tensorflow.python.distribute import strategy_combinations
|
|
from tensorflow.python.distribute import test_util
|
|
from tensorflow.python.distribute import values_v2
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.eager import test
|
|
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
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import resource_variable_ops
|
|
from tensorflow.python.ops import variables as variables_lib
|
|
|
|
|
|
class _VariableInterfaceTestBase(test.TestCase, parameterized.TestCase):
|
|
# This test verifies that DistributedVariable/AutoSyncVariable conforms to
|
|
# Variable and ResourceVariable interface, i.e. the methods and properties are
|
|
# all defined. It verifies methods and properties that have the same code path
|
|
# under different replicas/devices as well. It is not intended to verify
|
|
# methods and properties that behave differently under different
|
|
# replicas/devices; those should be covered separate tests.
|
|
|
|
def create_variable(self, initial_value=1., **kwargs):
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def devices(self):
|
|
return ["CPU:0", "CPU:1"]
|
|
|
|
# ==== Begin Variable interface ===
|
|
# Please follow the same order as methods and properties defined in
|
|
# tf.Variable.
|
|
|
|
def testStringify(self):
|
|
v = self.create_variable()
|
|
self.assertIsInstance(v.__str__(), str)
|
|
self.assertIsInstance(v.__repr__(), str)
|
|
|
|
def testDenseRead(self):
|
|
v = self.create_variable(1.)
|
|
self.assertEqual(v.value(), 1.)
|
|
self.assertEqual(v.read_value(), 1.)
|
|
|
|
def testShape(self):
|
|
v = self.create_variable([1.])
|
|
self.assertEqual(v.shape, (1,))
|
|
self.assertEqual(v.get_shape(), (1,))
|
|
v.set_shape((1,))
|
|
with self.assertRaisesRegex(ValueError, "not compatible"):
|
|
v.set_shape((1, 1))
|
|
|
|
@combinations.generate(combinations.combine(trainable=[True, False]))
|
|
def testTrainable(self, trainable):
|
|
v = self.create_variable(trainable=trainable)
|
|
self.assertEqual(v.trainable, trainable)
|
|
|
|
@combinations.generate(
|
|
combinations.combine(synchronization=[
|
|
variables_lib.VariableSynchronization.ON_READ,
|
|
variables_lib.VariableSynchronization.ON_WRITE,
|
|
variables_lib.VariableSynchronization.AUTO,
|
|
variables_lib.VariableSynchronization.NONE,
|
|
]))
|
|
def testSynchronization(self, synchronization):
|
|
v = self.create_variable(synchronization=synchronization)
|
|
self.assertEqual(v.synchronization, synchronization)
|
|
|
|
@combinations.generate(
|
|
combinations.combine(aggregation=[
|
|
variables_lib.VariableAggregation.MEAN,
|
|
variables_lib.VariableAggregation.SUM,
|
|
variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
|
|
variables_lib.VariableAggregation.NONE,
|
|
]))
|
|
def testAggregation(self, aggregation):
|
|
v = self.create_variable(aggregation=aggregation)
|
|
self.assertEqual(v.aggregation, aggregation)
|
|
|
|
@combinations.generate(combinations.combine(mode="graph"))
|
|
def testEval(self):
|
|
v = self.create_variable(1.)
|
|
with self.cached_session():
|
|
self.evaluate(variables_lib.global_variables_initializer())
|
|
self.assertEqual(v.eval(), 1.)
|
|
|
|
def testInitialValueEager(self):
|
|
v = self.create_variable(1.)
|
|
with self.assertRaises(RuntimeError):
|
|
v.initial_value # pylint: disable=pointless-statement
|
|
|
|
@combinations.generate(combinations.combine(mode="graph"))
|
|
def testInitialValueGraph(self):
|
|
v = self.create_variable(1.)
|
|
self.assertEqual(self.evaluate(v.initial_value), 1.)
|
|
|
|
def testConstraint(self):
|
|
v = self.create_variable(constraint=lambda x: x + 1.)
|
|
self.assertEqual(v.constraint(1.), 2.)
|
|
|
|
def testDenseUpdate(self):
|
|
v = self.create_variable(1.)
|
|
self.assertEqual(
|
|
v.assign(2., use_locking=True, name="assign", read_value=True), 2.)
|
|
self.assertIsNone(v.assign(3., read_value=False))
|
|
self.assertEqual(v, 3.)
|
|
self.assertEqual(
|
|
v.assign_add(1., use_locking=True, name="assign_add", read_value=True),
|
|
4.)
|
|
self.assertIsNone(v.assign_add(1., read_value=False))
|
|
self.assertEqual(v, 5.)
|
|
self.assertEqual(
|
|
v.assign_sub(1., use_locking=True, name="assign_sub", read_value=True),
|
|
4.)
|
|
self.assertIsNone(v.assign_sub(1., read_value=False))
|
|
self.assertEqual(v, 3.)
|
|
|
|
@def_function.function
|
|
def f():
|
|
self.assertIsInstance(v.assign(1., read_value=False), ops.Operation)
|
|
self.assertIsInstance(v.assign_add(1., read_value=False), ops.Operation)
|
|
self.assertIsInstance(v.assign_sub(1., read_value=False), ops.Operation)
|
|
|
|
f()
|
|
|
|
def testSparseUpdate(self):
|
|
v = self.create_variable([0., 0., 0.])
|
|
self.assertAllEqual(
|
|
v.scatter_add(
|
|
_make_index_slices(values=[1., 2.], indices=[0, 2]),
|
|
use_locking=True,
|
|
name="add"), [1., 0., 2.])
|
|
self.assertAllEqual(
|
|
v.scatter_div(
|
|
_make_index_slices(values=[4., 2.], indices=[0, 2]),
|
|
use_locking=True,
|
|
name="div"), [0.25, 0., 1.])
|
|
self.assertAllEqual(
|
|
v.scatter_max(
|
|
_make_index_slices(values=[1., 0.5], indices=[1, 2]),
|
|
use_locking=True,
|
|
name="max"), [0.25, 1., 1.])
|
|
self.assertAllEqual(
|
|
v.scatter_min(
|
|
_make_index_slices(values=[1., 0.5], indices=[0, 1]),
|
|
use_locking=True,
|
|
name="min"), [0.25, 0.5, 1.])
|
|
self.assertAllEqual(
|
|
v.scatter_mul(
|
|
_make_index_slices(values=[2., 0.5], indices=[0, 1]),
|
|
use_locking=True,
|
|
name="mul"), [0.5, 0.25, 1.])
|
|
self.assertAllEqual(
|
|
v.scatter_sub(
|
|
_make_index_slices(values=[2., 0.5], indices=[0, 1]),
|
|
use_locking=True,
|
|
name="sub"), [-1.5, -0.25, 1.])
|
|
self.assertAllEqual(
|
|
v.scatter_update(
|
|
_make_index_slices(values=[2., 0.5], indices=[0, 1]),
|
|
use_locking=True,
|
|
name="update"), [2., 0.5, 1.])
|
|
self.assertAllEqual(
|
|
v.batch_scatter_update(
|
|
_make_index_slices(values=[1., 1.5], indices=[0, 1]),
|
|
use_locking=True,
|
|
name="update"), [1., 1.5, 1.])
|
|
|
|
def testSparseNdUpdate(self):
|
|
v = self.create_variable([0., 0., 0., 0.])
|
|
self.assertAllEqual(
|
|
v.scatter_nd_sub([[3], [1]], [1., 2.], name="sub"), [0., -2., 0., -1.])
|
|
self.assertAllEqual(
|
|
v.scatter_nd_add([[2], [0]], [1., 2.], name="add"), [2., -2., 1., -1.])
|
|
self.assertAllEqual(
|
|
v.scatter_nd_update([[1], [3]], [3., 3.], name="update"),
|
|
[2., 3., 1., 3.])
|
|
|
|
def testSparseRead(self):
|
|
v = self.create_variable([[1., 2.], [3., 4.]])
|
|
self.assertAllEqual(
|
|
v.sparse_read([1, 0], name="read"), [[3., 4.], [1., 2.]])
|
|
self.assertAllEqual(
|
|
v.gather_nd([[1, 0], [0, 1]], name="gather_nd"), [3., 2.])
|
|
|
|
def testTensorConversion(self):
|
|
v = self.create_variable([1.])
|
|
self.assertEqual(ops.convert_to_tensor(v), [1.])
|
|
|
|
def testHash(self):
|
|
v = self.create_variable()
|
|
w = self.create_variable()
|
|
d = {}
|
|
with self.assertRaises(TypeError):
|
|
d[v] = 1
|
|
d[v.ref()] = 1
|
|
self.assertEqual(d[v.ref()], 1)
|
|
self.assertNotIn(w.ref(), d)
|
|
|
|
@combinations.generate(combinations.combine(mode="graph"))
|
|
def testHashGraph(self):
|
|
v = self.create_variable()
|
|
w = self.create_variable()
|
|
d = {v: 1}
|
|
self.assertEqual(d[v], 1)
|
|
self.assertNotIn(w, d)
|
|
|
|
def testEquality(self):
|
|
v = self.create_variable(1.)
|
|
w = self.create_variable(2.)
|
|
x = self.create_variable(1.)
|
|
self.assertEqual(v, x)
|
|
self.assertNotEqual(v, w)
|
|
|
|
@combinations.generate(combinations.combine(mode="graph"))
|
|
def testEqualityGraph(self):
|
|
# In legacy graph mode, tensor equality is object equality
|
|
v = self.create_variable(1.)
|
|
w = self.create_variable(1.)
|
|
self.assertNotEqual(v, w)
|
|
self.assertEqual(v, v)
|
|
|
|
def testIteration(self):
|
|
v = self.create_variable([1.])
|
|
self.assertEqual([1.], list(iter(v)))
|
|
|
|
def testProperties(self):
|
|
v = self.create_variable()
|
|
self.assertIsInstance(v.name, str)
|
|
# _shared_name is also part of the interface. E.g. it's used in optimizer to
|
|
# determine slot variable key.
|
|
self.assertIsInstance(v._shared_name, str)
|
|
self.assertIsNone(v.initializer)
|
|
self.assertIsInstance(v.device, str)
|
|
self.assertEqual(v.dtype, dtypes.float32)
|
|
with self.assertRaises(AttributeError):
|
|
v.op # pylint: disable=pointless-statement
|
|
with self.assertRaises(AttributeError):
|
|
v.graph # pylint: disable=pointless-statement
|
|
|
|
@combinations.generate(combinations.combine(mode="graph"))
|
|
def testPropertiesGraph(self):
|
|
v = self.create_variable()
|
|
self.assertIsInstance(v.initializer, ops.Operation)
|
|
self.assertIsInstance(v.op, ops.Operation)
|
|
self.assertIsInstance(v.graph, ops.Graph)
|
|
|
|
def testProtoConversion(self):
|
|
# to_proto and from_proto are not supported.
|
|
v = self.create_variable([1, 2])
|
|
with self.assertRaises(TypeError):
|
|
v.to_proto()
|
|
with self.assertRaises(TypeError):
|
|
v.from_proto(variable_def=None)
|
|
|
|
def testSaveSliceInfo(self):
|
|
v = self.create_variable()
|
|
slice_info = variables_lib.Variable.SaveSliceInfo()
|
|
v._set_save_slice_info(slice_info)
|
|
self.assertIs(v._get_save_slice_info(), slice_info)
|
|
# Some code accesses _save_slice_info directly without using the getter.
|
|
self.assertIs(v._save_slice_info, slice_info)
|
|
|
|
def testOperatorOverride(self):
|
|
v = self.create_variable(7)
|
|
self.assertEqual(v + 1, 8)
|
|
self.assertEqual(3 + v, 10)
|
|
self.assertEqual(v + v, 14)
|
|
self.assertEqual(v - 2, 5)
|
|
self.assertEqual(13 - v, 6)
|
|
self.assertEqual(v - v, 0)
|
|
self.assertEqual(v * 2, 14)
|
|
self.assertEqual(3 * v, 21)
|
|
self.assertEqual(v * v, 49)
|
|
self.assertEqual(v / 2, 3.5)
|
|
self.assertEqual(14 / v, 2.)
|
|
self.assertEqual(v // 2, 3)
|
|
self.assertEqual(15 // v, 2)
|
|
self.assertEqual(v % 2, 1)
|
|
self.assertEqual(16 % v, 2)
|
|
# pylint: disable=g-generic-assert
|
|
self.assertTrue(v < 12)
|
|
self.assertTrue(v <= 12)
|
|
self.assertFalse(v > 12)
|
|
self.assertFalse(v >= 12)
|
|
self.assertFalse(12 < v)
|
|
self.assertFalse(12 <= v)
|
|
self.assertTrue(12 > v)
|
|
self.assertTrue(12 >= v)
|
|
# pylint: enable=g-generic-assert
|
|
self.assertEqual(v & 3, 3)
|
|
self.assertEqual(11 & v, 3)
|
|
self.assertEqual(v | 8, 15)
|
|
self.assertEqual(16 | v, 23)
|
|
self.assertEqual(v ^ 3, 4)
|
|
self.assertEqual(11 ^ v, 12)
|
|
self.assertEqual(pow(v, 3), 343)
|
|
# TODO(b/178748613): pow(v, 3, 10) fails.
|
|
self.assertEqual(pow(2, v), 128)
|
|
self.assertEqual(-v, -7)
|
|
self.assertEqual(~v, ~7)
|
|
self.assertEqual(abs(v), 7)
|
|
|
|
def testSlice(self):
|
|
v = self.create_variable([1., 2., 3.])
|
|
self.assertEqual(v[1], 2.)
|
|
v[2].assign(4.)
|
|
self.assertAllEqual(v, [1., 2., 4.])
|
|
|
|
# ==== End Variable interface ===
|
|
|
|
# ==== Begin ResourceVariable interface ===
|
|
def testHandle(self):
|
|
v = self.create_variable()
|
|
self.assertIsInstance(v.handle, tensor.Tensor)
|
|
self.assertEqual(v.handle.dtype, dtypes.resource)
|
|
|
|
def testInGraphMode(self):
|
|
# This is protected but used in a lot of places internally.
|
|
v = self.create_variable()
|
|
self.assertFalse(v._in_graph_mode)
|
|
|
|
def testUniqueId(self):
|
|
# This is used in optimizer as part of slot variable key.
|
|
v = self.create_variable()
|
|
w = self.create_variable()
|
|
self.assertNotEqual(v._unique_id, w._unique_id)
|
|
|
|
def testIsResourceVariable(self):
|
|
v = self.create_variable()
|
|
self.assertTrue(resource_variable_ops.is_resource_variable(v))
|
|
# ==== End ResourceVariable interface ===
|
|
|
|
@combinations.generate(combinations.combine(mode="graph"))
|
|
def testAsGraphElement(self):
|
|
g = ops.Graph()
|
|
with g.as_default():
|
|
v = self.create_variable(1.)
|
|
g.finalize()
|
|
self.evaluate(v.initializer)
|
|
# _as_graph_element shouldn't create new operations.
|
|
self.assertEqual(self.evaluate(v._as_graph_element()), 1.)
|
|
|
|
|
|
class DistributedVariableInterfaceTest(_VariableInterfaceTestBase):
|
|
|
|
def create_variable(self, initial_value=1., **kwargs):
|
|
variables = []
|
|
for device in self.devices:
|
|
with ops.device(device):
|
|
variables.append(
|
|
variables_lib.Variable(initial_value, **kwargs))
|
|
return values_v2.DistributedVariable(variables)
|
|
|
|
|
|
# Prevent the base class from running.
|
|
del _VariableInterfaceTestBase
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
strategy=[
|
|
strategy_combinations.tpu_strategy,
|
|
strategy_combinations.mirrored_strategy_with_two_cpus,
|
|
strategy_combinations.mirrored_strategy_with_two_gpus,
|
|
],
|
|
enable_packed_handle=[True, False],
|
|
tf_function=[combinations.tf_function, combinations.no_tf_function]))
|
|
class DistributedVariableTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def create_variable(self, strategy, initial_value, enable_packed_handle,
|
|
**kwargs):
|
|
variables = []
|
|
for device in strategy.extended.parameter_devices:
|
|
with ops.device(device):
|
|
variables.append(variables_lib.Variable(initial_value, **kwargs))
|
|
return values_v2.DistributedVariable(
|
|
variables, enable_packed_handle=enable_packed_handle)
|
|
|
|
def assertReplica(self, distributed_var, values):
|
|
for var, value in zip(distributed_var._variables, values):
|
|
self.assertAllEqual(var, value)
|
|
|
|
def testRead(self, strategy, enable_packed_handle, tf_function):
|
|
v = self.create_variable(strategy, 0., enable_packed_handle)
|
|
|
|
with ops.device(strategy.extended.parameter_devices[0]):
|
|
v.assign(1.)
|
|
with ops.device(strategy.extended.parameter_devices[1]):
|
|
v.assign(2.)
|
|
|
|
@tf_function
|
|
def read_device0():
|
|
with ops.device(strategy.extended.parameter_devices[0]):
|
|
return v.read_value(), v.value()
|
|
|
|
@tf_function
|
|
def read_device1():
|
|
with ops.device(strategy.extended.parameter_devices[1]):
|
|
return v.read_value(), v.value()
|
|
|
|
@tf_function
|
|
def read_other_device():
|
|
with ops.device("CPU:0"):
|
|
return v.read_value(), v.value()
|
|
|
|
self.assertAllEqual(read_device0(), [1., 1.])
|
|
self.assertAllEqual(read_device1(), [2., 2.])
|
|
self.assertAllEqual(read_other_device(), [1., 1.])
|
|
|
|
def testAssign(self, strategy, enable_packed_handle, tf_function):
|
|
v = self.create_variable(strategy, 0., enable_packed_handle)
|
|
|
|
@tf_function
|
|
def update_device0():
|
|
with ops.device(strategy.extended.parameter_devices[0]):
|
|
v.assign(1.)
|
|
|
|
@tf_function
|
|
def update_device1():
|
|
with ops.device(strategy.extended.parameter_devices[1]):
|
|
v.assign(2.)
|
|
|
|
update_device0()
|
|
update_device1()
|
|
self.assertReplica(v, [1., 2.])
|
|
|
|
with ops.device("CPU:0"):
|
|
# Update the primary replica.
|
|
v.assign(3.)
|
|
self.assertReplica(v, [3., 2.])
|
|
|
|
def testStrategyRun(self, strategy, enable_packed_handle, tf_function):
|
|
if (test_util.is_tpu_strategy(strategy) and
|
|
tf_function is combinations.no_tf_function):
|
|
self.skipTest("tpu doesn't support eager")
|
|
v = self.create_variable(strategy, 0., enable_packed_handle)
|
|
|
|
@tf_function
|
|
def update(per_replica):
|
|
v.assign(per_replica)
|
|
|
|
@tf_function
|
|
def read():
|
|
return v.read_value()
|
|
|
|
strategy.run(
|
|
update, args=(test_util.create_per_replica(strategy, [1., 2.]),))
|
|
self.assertReplica(v, [1., 2.])
|
|
self.assertAllEqual(
|
|
test_util.gather(strategy, strategy.run(read)), [1., 2.])
|
|
|
|
|
|
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__":
|
|
test_util.main()
|