wehub-resource-sync
765 lines
28 KiB
Python
765 lines
28 KiB
Python
# Copyright 2020 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 common methods in strategy classes."""
|
|
|
|
from absl.testing import parameterized
|
|
|
|
from tensorflow.python.data.ops import dataset_ops
|
|
from tensorflow.python.distribute import combinations
|
|
from tensorflow.python.distribute import distribute_lib
|
|
from tensorflow.python.distribute import multi_worker_test_base
|
|
from tensorflow.python.distribute import reduce_util
|
|
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.distribute.collective_all_reduce_strategy import CollectiveAllReduceStrategy
|
|
from tensorflow.python.eager import def_function
|
|
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.ops import array_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.platform import test
|
|
from tensorflow.python.util import nest
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
strategy=[
|
|
strategy_combinations.multi_worker_mirrored_2x1_cpu,
|
|
strategy_combinations.multi_worker_mirrored_2x1_gpu,
|
|
] + strategy_combinations.all_strategies,
|
|
mode=['eager']))
|
|
class StrategyTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def testCaptureReplicaId(self, strategy):
|
|
m = {}
|
|
|
|
@def_function.function
|
|
def f():
|
|
return distribute_lib.get_replica_context().replica_id_in_sync_group
|
|
|
|
@def_function.function
|
|
def g():
|
|
# Make g() a stateful function so it's traced twice.
|
|
if m.get('v', None) is None:
|
|
m['v'] = variables.Variable(0.)
|
|
return strategy.run(f)
|
|
|
|
g()
|
|
|
|
def testMergeCallInitScope(self, strategy):
|
|
with strategy.scope():
|
|
|
|
@def_function.function
|
|
def fn():
|
|
|
|
def merge_fn(unused_strat):
|
|
|
|
y = constant_op.constant(11)
|
|
return y
|
|
|
|
def replica_fn():
|
|
|
|
with ops.init_scope():
|
|
y = distribute_lib.get_replica_context().merge_call(merge_fn)
|
|
z = y + 1
|
|
return z
|
|
|
|
return strategy.run(replica_fn)
|
|
|
|
result = strategy.experimental_local_results(fn())
|
|
self.assertAllClose(result, [12] * _get_num_replicas_per_client(strategy))
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
distribution=[
|
|
strategy_combinations.mirrored_strategy_with_cpu_1_and_2,
|
|
strategy_combinations.multi_worker_mirrored_2x2_gpu,
|
|
strategy_combinations.tpu_strategy
|
|
],
|
|
mode=['graph', 'eager']))
|
|
class StrategyLocalResultTest(test.TestCase):
|
|
|
|
def testLocalResultForDictionary(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
return {'a': constant_op.constant(1.), 'b': constant_op.constant(2.)}
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
got = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertEqual(got, ({'a': 1., 'b': 2.}, {'a': 1., 'b': 2.}))
|
|
|
|
def testLocalResultForList(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
return [constant_op.constant(1.), constant_op.constant(2.)]
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
got = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertEqual(got, ([1., 2.], [1., 2.]))
|
|
|
|
def testLocalResultForTuple(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
return (constant_op.constant(1.), constant_op.constant(2.),
|
|
constant_op.constant(3.))
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
got = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertEqual(got, ((1., 2., 3.), (1., 2., 3.)))
|
|
|
|
def testLocalResultForNestedStruct(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
return ({
|
|
'a': constant_op.constant(1.),
|
|
'b': constant_op.constant(2.)
|
|
}, {
|
|
'a': constant_op.constant(4.),
|
|
'b': constant_op.constant(6.)
|
|
})
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
got = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertEqual(got, (({
|
|
'a': 1.,
|
|
'b': 2.
|
|
}, {
|
|
'a': 4.,
|
|
'b': 6.
|
|
}), ({
|
|
'a': 1.,
|
|
'b': 2.
|
|
}, {
|
|
'a': 4.,
|
|
'b': 6.
|
|
})))
|
|
|
|
def testLocalResultForNestedStructWithoutTensor(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
return {'a': 1., 'b': 2.}
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
v = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertIsInstance(v, tuple)
|
|
self.assertAllEqual(v, ({'a': 1., 'b': 2.}, {'a': 1., 'b': 2.}))
|
|
|
|
def testLocalResultForScalarValue(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
return distribution.extended._get_local_replica_id(
|
|
distribute_lib.get_replica_context().replica_id_in_sync_group)
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
v = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertIsInstance(v, tuple)
|
|
self.assertEqual(v, (0, 1))
|
|
|
|
def testLocalResultForDictionaryDifferentReplicas(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
replica_id = distribution.extended._get_local_replica_id(
|
|
distribute_lib.get_replica_context().replica_id_in_sync_group)
|
|
return {
|
|
'a': math_ops.cast(replica_id + 1, dtype=float),
|
|
'b': math_ops.cast(replica_id + 2, dtype=float)
|
|
}
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
got = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertAllEqual(got, ({'a': 1., 'b': 2.}, {'a': 2., 'b': 3.}))
|
|
|
|
def testLocalResultForTensor(self, distribution):
|
|
|
|
@def_function.function
|
|
def model_fn():
|
|
return constant_op.constant([2., 3.])
|
|
|
|
with distribution.scope():
|
|
result = distribution.run(model_fn)
|
|
v = self.evaluate(distribution.experimental_local_results(result))
|
|
self.assertAllEqual(v, ([2., 3.], [2., 3.]))
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
strategy=[
|
|
strategy_combinations.multi_worker_mirrored_2x1_cpu,
|
|
strategy_combinations.multi_worker_mirrored_2x1_gpu,
|
|
] + strategy_combinations.all_strategies,
|
|
mode=['eager']))
|
|
class ReduceTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def testBasic(self, strategy):
|
|
per_replica_value = strategy.experimental_distribute_values_from_function(
|
|
lambda _: array_ops.ones((), dtypes.float32))
|
|
|
|
def fn_eager():
|
|
|
|
return strategy.reduce(
|
|
reduce_util.ReduceOp.SUM, value=per_replica_value, axis=None)
|
|
|
|
fn_graph = def_function.function(fn_eager)
|
|
# Run reduce under the strategy scope to explicitly enter
|
|
# strategy default_device scope.
|
|
with strategy.scope():
|
|
self.assertEqual(fn_eager().numpy(), 1.0 * strategy.num_replicas_in_sync)
|
|
self.assertEqual(fn_graph().numpy(), 1.0 * strategy.num_replicas_in_sync)
|
|
|
|
# Run reduce without a strategy scope to implicitly enter
|
|
# strategy default_device scope.
|
|
self.assertEqual(fn_eager().numpy(), 1.0 * strategy.num_replicas_in_sync)
|
|
self.assertEqual(fn_graph().numpy(), 1.0 * strategy.num_replicas_in_sync)
|
|
|
|
def testAxis(self, strategy):
|
|
|
|
@def_function.function
|
|
def fn():
|
|
return constant_op.constant([1., 2.])
|
|
|
|
x = strategy.run(fn)
|
|
|
|
x_m = strategy.reduce(reduce_util.ReduceOp.MEAN, x, axis=0)
|
|
self.assertEqual(1.5, x_m)
|
|
x_s = strategy.reduce(reduce_util.ReduceOp.SUM, x, axis=0)
|
|
self.assertEqual(3 * strategy.num_replicas_in_sync, x_s)
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
strategy=[
|
|
strategy_combinations.default_strategy,
|
|
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.multi_worker_mirrored_2x1_cpu,
|
|
strategy_combinations.multi_worker_mirrored_2x2_gpu,
|
|
strategy_combinations.multi_worker_mirrored_2x2_gpu_no_merge_call,
|
|
],
|
|
update_fn=['assign', 'assign_add', 'assign_sub'],
|
|
tf_function=[True, False],
|
|
mode=['eager']))
|
|
class ReplicaCtxUpdateTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def testDenseUpdate(self, strategy, tf_function, update_fn):
|
|
if strategy_test_lib.is_tpu_strategy(strategy) and (not tf_function):
|
|
self.skipTest('Skip TPUStrategy + eager combination.')
|
|
with strategy.scope():
|
|
distributed_variable1 = variables.Variable(5.0)
|
|
|
|
def replica_fn():
|
|
value = array_ops.constant(2.)
|
|
python_literal = 1.
|
|
replica_context = distribute_lib.get_replica_context()
|
|
fn_sets = {
|
|
'assign': lambda var, value: var.assign(value),
|
|
'assign_add': lambda var, value: var.assign_add(value),
|
|
'assign_sub': lambda var, value: var.assign_sub(value),
|
|
}
|
|
replica_context._update(
|
|
distributed_variable1, fn_sets[update_fn], args=(value,))
|
|
replica_context._update(
|
|
distributed_variable1, fn_sets[update_fn], args=(python_literal,))
|
|
|
|
if tf_function:
|
|
replica_fn = def_function.function(replica_fn)
|
|
strategy.run(replica_fn)
|
|
|
|
expected_result = {'assign': 1., 'assign_add': 8., 'assign_sub': 2.}
|
|
self.assertAllEqual(
|
|
strategy.experimental_local_results(distributed_variable1),
|
|
[expected_result[update_fn]] * _get_num_replicas_per_client(strategy))
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
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.tpu_strategy,
|
|
] + strategy_combinations.strategies_minus_tpu,
|
|
tf_function=[combinations.tf_function, combinations.no_tf_function],
|
|
mode=['eager']))
|
|
class ReplicaCtxAllReduceTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def testDense(self, strategy, tf_function):
|
|
if (strategy_test_lib.is_tpu_strategy(strategy) and
|
|
tf_function is combinations.no_tf_function):
|
|
self.skipTest('Skip TPUStrategy + eager combination.')
|
|
|
|
@tf_function
|
|
def fn():
|
|
|
|
def replica_fn():
|
|
value = array_ops.identity(1.0)
|
|
reduced = strategy.extended._replica_ctx_all_reduce(
|
|
reduce_util.ReduceOp.SUM, value)
|
|
return reduced
|
|
|
|
return strategy.experimental_local_results(strategy.run(replica_fn))
|
|
|
|
got = fn()[0]
|
|
self.assertEqual(got, 1.0 * strategy.num_replicas_in_sync)
|
|
|
|
def testSparse(self, strategy, tf_function):
|
|
if tf_function is combinations.no_tf_function:
|
|
self.skipTest('Skip IndexedSlices + eager combination.')
|
|
|
|
@tf_function
|
|
def fn():
|
|
|
|
def replica_fn():
|
|
value = indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1]))
|
|
reduced = strategy.extended._replica_ctx_all_reduce(
|
|
reduce_util.ReduceOp.SUM, value)
|
|
return reduced
|
|
|
|
return strategy.experimental_local_results(strategy.run(replica_fn))
|
|
|
|
got = fn()[0]
|
|
expect = indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0 * strategy.num_replicas_in_sync]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1]))
|
|
self.assertAllEqual(
|
|
ops.convert_to_tensor(got), ops.convert_to_tensor(expect))
|
|
|
|
def testNestedInput(self, strategy, tf_function):
|
|
if tf_function is combinations.no_tf_function:
|
|
self.skipTest('Skip IndexedSlices + eager combination.')
|
|
|
|
@tf_function
|
|
def fn():
|
|
|
|
def replica_fn():
|
|
value = (array_ops.identity(1.0),
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1])),
|
|
array_ops.identity(2.0),
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[2.0]]),
|
|
indices=array_ops.identity([1]),
|
|
dense_shape=array_ops.identity([5, 1])))
|
|
reduced = strategy.extended._replica_ctx_all_reduce(
|
|
reduce_util.ReduceOp.SUM, value)
|
|
return reduced
|
|
|
|
return strategy.experimental_local_results(strategy.run(replica_fn))
|
|
|
|
got = fn()[0]
|
|
expect = (1.0 * strategy.num_replicas_in_sync,
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity(
|
|
[[1.0 * strategy.num_replicas_in_sync]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1])),
|
|
2.0 * strategy.num_replicas_in_sync,
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity(
|
|
[[2.0 * strategy.num_replicas_in_sync]]),
|
|
indices=array_ops.identity([1]),
|
|
dense_shape=array_ops.identity([5, 1])))
|
|
|
|
self.assertAllClose(
|
|
nest.map_structure(ops.convert_to_tensor, got),
|
|
nest.map_structure(ops.convert_to_tensor, expect))
|
|
|
|
def testSyncOnReadVariableInput(self, strategy, tf_function):
|
|
if (not strategy_test_lib.is_mirrored_strategy(strategy) and
|
|
not strategy_test_lib.is_multi_worker_mirrored_strategy(strategy) and
|
|
not strategy_test_lib.is_tpu_strategy(strategy)):
|
|
self.skipTest('Skip strategies not using SyncOnReadVariables.')
|
|
if (strategy_test_lib.is_tpu_strategy(strategy) and
|
|
tf_function is combinations.no_tf_function):
|
|
self.skipTest('Skip TPUStrategy + eager combination.')
|
|
if (strategy_test_lib.is_multi_worker_mirrored_strategy(strategy) and
|
|
tf_function is combinations.tf_function):
|
|
self.skipTest('Skip MWMS + graph combination until b/228512201 is fixed.')
|
|
|
|
with strategy.scope():
|
|
var = variables.Variable(
|
|
0.0,
|
|
synchronization=variables.VariableSynchronization.ON_READ,
|
|
aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA)
|
|
|
|
@tf_function
|
|
def replica_fn():
|
|
replica_context = distribute_lib.get_replica_context()
|
|
replica_id = replica_context.replica_id_in_sync_group
|
|
var.assign(math_ops.cast(replica_id, dtype=float) * 3.0)
|
|
|
|
return replica_context.all_reduce(reduce_util.ReduceOp.SUM, var)
|
|
|
|
if strategy_test_lib.is_multi_worker_mirrored_strategy(strategy):
|
|
client_local_replica_num = strategy.extended._num_devices_per_worker
|
|
else:
|
|
client_local_replica_num = strategy.num_replicas_in_sync
|
|
|
|
workers_num = strategy.num_replicas_in_sync
|
|
expected_sum = sum(range(workers_num)) * 3.0
|
|
|
|
# Expand the values on each replica if multiple devices are used; otherwise
|
|
# simple read the value of the Tensor.
|
|
result = strategy.run(replica_fn)
|
|
if hasattr(result, 'values'):
|
|
result = result.values
|
|
result = nest.flatten(result)
|
|
|
|
# Iterate through all replicas and verify the reduce sum result.
|
|
for i in range(client_local_replica_num):
|
|
self.assertEqual(result[i].numpy(), expected_sum)
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
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.tpu_strategy,
|
|
] + strategy_combinations.strategies_minus_tpu,
|
|
tf_function=[combinations.tf_function, combinations.no_tf_function],
|
|
mode=['eager']))
|
|
class AllReduceTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def testDense(self, strategy, tf_function):
|
|
if (strategy_test_lib.is_tpu_strategy(strategy) and
|
|
tf_function is combinations.no_tf_function):
|
|
self.skipTest('Skip TPUStrategy + eager combination.')
|
|
|
|
@tf_function
|
|
def fn():
|
|
|
|
def replica_fn():
|
|
value = array_ops.identity(1.0)
|
|
rep_ctx = distribute_lib.get_replica_context()
|
|
reduced = rep_ctx.all_reduce(reduce_util.ReduceOp.SUM, value)
|
|
return reduced
|
|
|
|
return strategy.experimental_local_results(strategy.run(replica_fn))
|
|
|
|
got = fn()[0]
|
|
self.assertEqual(got, 1.0 * strategy.num_replicas_in_sync)
|
|
|
|
def testSparse(self, strategy, tf_function):
|
|
if tf_function is combinations.no_tf_function:
|
|
self.skipTest('Skip IndexedSlices + eager combination.')
|
|
|
|
@tf_function
|
|
def fn():
|
|
|
|
def replica_fn():
|
|
value = indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1]))
|
|
rep_ctx = distribute_lib.get_replica_context()
|
|
reduced = rep_ctx.all_reduce(reduce_util.ReduceOp.MEAN, value)
|
|
return reduced
|
|
|
|
return strategy.experimental_local_results(strategy.run(replica_fn))
|
|
|
|
got = fn()[0]
|
|
|
|
if not strategy_test_lib.is_tpu_strategy(strategy):
|
|
self.assertIsInstance(got, indexed_slices.IndexedSlices)
|
|
expect = indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1]))
|
|
self.assertAllEqual(
|
|
ops.convert_to_tensor(got), ops.convert_to_tensor(expect))
|
|
|
|
def testSparseTuple(self, strategy, tf_function):
|
|
if tf_function is combinations.no_tf_function:
|
|
self.skipTest('Skip IndexedSlices + eager combination.')
|
|
|
|
@tf_function
|
|
def fn():
|
|
|
|
def replica_fn():
|
|
value1 = indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1]))
|
|
value2 = indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[2.0]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1]))
|
|
rep_ctx = distribute_lib.get_replica_context()
|
|
reduced = rep_ctx.all_reduce(reduce_util.ReduceOp.SUM, [value1, value2])
|
|
return reduced
|
|
|
|
return strategy.experimental_local_results(strategy.run(replica_fn))
|
|
|
|
got = fn()[0]
|
|
|
|
if not strategy_test_lib.is_tpu_strategy(strategy):
|
|
for g in got:
|
|
self.assertIsInstance(g, indexed_slices.IndexedSlices)
|
|
expect = [
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0 * strategy.num_replicas_in_sync]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1])),
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[2.0 * strategy.num_replicas_in_sync]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1]))
|
|
]
|
|
self.assertAllEqual(
|
|
nest.map_structure(ops.convert_to_tensor, got),
|
|
nest.map_structure(ops.convert_to_tensor, expect))
|
|
|
|
def testNestedInput(self, strategy, tf_function):
|
|
if tf_function is combinations.no_tf_function:
|
|
self.skipTest('Skip IndexedSlices + eager combination.')
|
|
|
|
@tf_function
|
|
def fn():
|
|
|
|
def replica_fn():
|
|
value = (array_ops.identity(1.0),
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[1.0]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1])),
|
|
array_ops.identity(2.0),
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity([[2.0]]),
|
|
indices=array_ops.identity([1]),
|
|
dense_shape=array_ops.identity([5, 1])))
|
|
rep_ctx = distribute_lib.get_replica_context()
|
|
reduced = rep_ctx.all_reduce(reduce_util.ReduceOp.SUM, value)
|
|
return reduced
|
|
|
|
return strategy.experimental_local_results(strategy.run(replica_fn))
|
|
|
|
got = fn()[0]
|
|
expect = (1.0 * strategy.num_replicas_in_sync,
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity(
|
|
[[1.0 * strategy.num_replicas_in_sync]]),
|
|
indices=array_ops.identity([0]),
|
|
dense_shape=array_ops.identity([5, 1])),
|
|
2.0 * strategy.num_replicas_in_sync,
|
|
indexed_slices.IndexedSlices(
|
|
values=array_ops.identity(
|
|
[[2.0 * strategy.num_replicas_in_sync]]),
|
|
indices=array_ops.identity([1]),
|
|
dense_shape=array_ops.identity([5, 1])))
|
|
|
|
self.assertAllClose(
|
|
nest.map_structure(ops.convert_to_tensor, got),
|
|
nest.map_structure(ops.convert_to_tensor, expect))
|
|
|
|
|
|
def _make_indexed_slices(values, indices, dense_shape):
|
|
tensor = indexed_slices.IndexedSlices(
|
|
values=constant_op.constant(values),
|
|
indices=constant_op.constant(indices),
|
|
dense_shape=constant_op.constant(dense_shape))
|
|
return tensor
|
|
|
|
|
|
def _get_num_replicas_per_client(strategy):
|
|
if isinstance(strategy, CollectiveAllReduceStrategy):
|
|
resolver = strategy.cluster_resolver
|
|
return max(nest.flatten(resolver.num_accelerators())[0], 1)
|
|
else:
|
|
return strategy.num_replicas_in_sync
|
|
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
strategy=[
|
|
strategy_combinations.multi_worker_mirrored_2x1_cpu,
|
|
strategy_combinations.multi_worker_mirrored_2x1_gpu,
|
|
],
|
|
mode=['eager']))
|
|
class DistributedCollectiveAllReduceStrategyTest(
|
|
strategy_test_lib.DistributionTestBase,
|
|
parameterized.TestCase):
|
|
|
|
def testDatasetFromFunction(self, strategy):
|
|
def dataset_fn(input_context):
|
|
global_batch_size = 10
|
|
batch_size = input_context.get_per_replica_batch_size(global_batch_size)
|
|
d = dataset_ops.DatasetV2.range(100).repeat().batch(batch_size)
|
|
return d.shard(input_context.num_input_pipelines,
|
|
input_context.input_pipeline_id)
|
|
|
|
expected_sum_on_workers = {'chief': 10, 'worker': 35}
|
|
input_iterator = iter(
|
|
strategy.distribute_datasets_from_function(dataset_fn))
|
|
|
|
@def_function.function
|
|
def run(iterator):
|
|
return strategy.experimental_local_results(iterator.get_next())
|
|
|
|
result = run(input_iterator)
|
|
sum_value = math_ops.reduce_sum(result)
|
|
self.assertEqual(
|
|
sum_value.numpy(),
|
|
expected_sum_on_workers[multi_worker_test_base.get_task_type()])
|
|
|
|
def testSimpleInputFromDatasetLastPartialBatch(self, strategy):
|
|
global_batch_size = 8
|
|
dataset = dataset_ops.DatasetV2.range(14).batch(
|
|
global_batch_size, drop_remainder=False)
|
|
input_iterator = iter(strategy.experimental_distribute_dataset(dataset))
|
|
|
|
@def_function.function
|
|
def run(input_iterator):
|
|
return strategy.run(lambda x: x, args=(next(input_iterator),))
|
|
|
|
# Let the complete batch go.
|
|
run(input_iterator)
|
|
|
|
# `result` is an incomplete batch
|
|
result = run(input_iterator)
|
|
expected_data_on_workers = {'chief': [8, 9, 10], 'worker': [11, 12, 13]}
|
|
self.assertAllEqual(
|
|
expected_data_on_workers[multi_worker_test_base.get_task_type()],
|
|
result.numpy(),
|
|
)
|
|
|
|
def testSimpleInputFromFnLastPartialBatch(self, strategy):
|
|
|
|
def dataset_fn(input_context):
|
|
global_batch_size = 8
|
|
batch_size = input_context.get_per_replica_batch_size(global_batch_size)
|
|
dataset = dataset_ops.DatasetV2.range(14).batch(
|
|
batch_size, drop_remainder=False)
|
|
return dataset.shard(input_context.num_input_pipelines,
|
|
input_context.input_pipeline_id)
|
|
|
|
input_iterator = iter(
|
|
strategy.distribute_datasets_from_function(dataset_fn))
|
|
|
|
@def_function.function
|
|
def run(input_iterator):
|
|
return strategy.run(lambda x: x, args=(next(input_iterator),))
|
|
|
|
# Let the complete batch go.
|
|
run(input_iterator)
|
|
# `result` is an incomplete batch
|
|
result = run(input_iterator)
|
|
|
|
expected_data_on_worker = {'chief': [8, 9, 10, 11], 'worker': [12, 13]}
|
|
self.assertAllEqual(
|
|
expected_data_on_worker[multi_worker_test_base.get_task_type()],
|
|
result.numpy())
|
|
|
|
def testReduceHostTensor(self, strategy):
|
|
reduced = strategy.reduce(
|
|
reduce_util.ReduceOp.SUM, array_ops.identity(1.), axis=None)
|
|
self.assertEqual(reduced.numpy(), 2.)
|
|
|
|
def testReduceToHostTensor(self, strategy):
|
|
value = array_ops.identity(1.)
|
|
reduced = strategy.extended.reduce_to(reduce_util.ReduceOp.SUM, value,
|
|
value)
|
|
self.assertEqual(reduced.numpy(), 2.)
|
|
|
|
def testBatchReduceToHostTensor(self, strategy):
|
|
value = array_ops.identity(1.)
|
|
reduced = strategy.extended.batch_reduce_to(reduce_util.ReduceOp.SUM,
|
|
[(value, value),
|
|
(value, value)])
|
|
self.assertAllEqual([2., 2.], reduced)
|
|
|
|
def testReduceDeviceTensors(self, strategy):
|
|
value = strategy.run(lambda: array_ops.identity(1.))
|
|
reduced = strategy.reduce(reduce_util.ReduceOp.SUM, value, axis=None)
|
|
self.assertEqual(reduced.numpy(), 2.)
|
|
|
|
def testReduceToDeviceTensors(self, strategy):
|
|
value = strategy.run(lambda: array_ops.identity(1.))
|
|
reduced = strategy.extended.reduce_to(reduce_util.ReduceOp.SUM, value,
|
|
value)
|
|
self.assertEqual(reduced.numpy(), 2.)
|
|
|
|
def testBatchReduceToDeviceTensors(self, strategy):
|
|
value = strategy.run(lambda: array_ops.identity(1.))
|
|
reduced = strategy.extended.batch_reduce_to(reduce_util.ReduceOp.SUM,
|
|
[(value, value),
|
|
(value, value)])
|
|
self.assertAllEqual([2., 2.], reduced)
|
|
|
|
# TODO(crccw): add a test that mixes device and host tensors after multi
|
|
# worker strategy combinations can run on a fixed number of GPUs.
|
|
|
|
|
|
class StrategyClusterResolverTest(test.TestCase, parameterized.TestCase):
|
|
|
|
@combinations.generate(
|
|
combinations.combine(
|
|
strategy=[strategy_combinations.multi_worker_mirrored_2x1_cpu] +
|
|
strategy_combinations.all_strategies,
|
|
mode=['eager']))
|
|
def testClusterResolverProperty(self, strategy):
|
|
# CollectiveAllReduceStrategy and TPUStrategy must have a cluster resolver.
|
|
# `None` otherwise.
|
|
resolver = strategy.cluster_resolver
|
|
if (not isinstance(strategy, CollectiveAllReduceStrategy) and
|
|
not strategy_test_lib.is_tpu_strategy(strategy)):
|
|
self.assertIsNone(resolver)
|
|
return
|
|
|
|
with strategy.scope():
|
|
self.assertIs(strategy.cluster_resolver, resolver)
|
|
|
|
self.assertTrue(hasattr(resolver, 'cluster_spec'))
|
|
self.assertTrue(hasattr(resolver, 'master'))
|
|
self.assertTrue(hasattr(resolver, 'num_accelerators'))
|
|
self.assertTrue(hasattr(resolver, 'task_id'))
|
|
self.assertTrue(hasattr(resolver, 'task_type'))
|
|
if isinstance(strategy, CollectiveAllReduceStrategy):
|
|
self.assertEqual(resolver.task_id, 0)
|
|
self.assertAllInSet(resolver.task_type, ['chief', 'worker'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_util.main()
|