Remove run_deprecated_v1 annotations from collective ops tests.
PiperOrigin-RevId: 317365063 Change-Id: Ibf13ad8629947becd40038d41ee213d3466b6292
This commit is contained in:
Ayush Dubey
TensorFlower Gardener
parent
3ae2cf9610
commit
c575e2ba93
@ -104,39 +104,42 @@ class CollectiveOpTest(test.TestCase):
|
||||
for i in range(group_size * num_instances):
|
||||
self.assertAllClose(results[i], expected, rtol=1e-5, atol=1e-5)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveReduce(self):
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=True)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=True)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveAutoGraphKey(self):
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=False)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=False)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testFp16Reduce(self):
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=True,
|
||||
fp16=True)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=True,
|
||||
fp16=True)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveMultipleConcurrentReduce(self):
|
||||
self._testMultipleConcurrentCollectiveReduce(
|
||||
[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
|
||||
[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2])
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testMultipleConcurrentCollectiveReduce(
|
||||
[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
|
||||
[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2])
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveTimeoutV1(self):
|
||||
timeout = 4.5
|
||||
kwargs = dict(
|
||||
@ -145,14 +148,17 @@ class CollectiveOpTest(test.TestCase):
|
||||
set_graph_key=True,
|
||||
timeout=timeout)
|
||||
|
||||
self._testCollectiveReduce(**kwargs)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(**kwargs)
|
||||
|
||||
start_time = time.time()
|
||||
with self.assertRaisesRegex(
|
||||
errors.DeadlineExceededError,
|
||||
'Collective has timed out waiting for other workers'):
|
||||
self._testCollectiveReduce(
|
||||
reported_group_size=len(kwargs['inputs']) + 1, **kwargs)
|
||||
with ops.Graph().as_default():
|
||||
with self.assertRaisesRegex(
|
||||
errors.DeadlineExceededError,
|
||||
'Collective has timed out waiting for other workers'):
|
||||
self._testCollectiveReduce(
|
||||
reported_group_size=len(kwargs['inputs']) + 1, **kwargs)
|
||||
elapsed = time.time() - start_time
|
||||
self.assertAllGreaterEqual(elapsed, timeout)
|
||||
|
||||
@ -199,17 +205,18 @@ class CollectiveOpTest(test.TestCase):
|
||||
elapsed = time.time() - start_time
|
||||
self.assertAllGreaterEqual(elapsed, timeout)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testNcclHintFallbackToRingReduce(self):
|
||||
"""Tests that setting `communication_hint=nccl` works on non-GPU builds."""
|
||||
if kernels.get_registered_kernels_for_op('NcclAllReduce'):
|
||||
self.skipTest('Run only on non-GPU environments')
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=False,
|
||||
communication_hint='nccl')
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
|
||||
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
|
||||
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
|
||||
set_graph_key=False,
|
||||
communication_hint='nccl')
|
||||
|
||||
def _testWhile(self, num_vars, num_iterations, key_base):
|
||||
group_size = 2
|
||||
@ -262,15 +269,16 @@ class CollectiveOpTest(test.TestCase):
|
||||
[((1 << (num_iterations + v)) * 1.) for v in range(num_vars)]
|
||||
for _ in range(group_size)])
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testSimpleWhile(self):
|
||||
self._testWhile(num_vars=1, num_iterations=4, key_base=20)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testWhile(num_vars=1, num_iterations=4, key_base=20)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testWhileMultipleAllReduce(self):
|
||||
self._testWhile(num_vars=2, num_iterations=4, key_base=20)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testWhile(num_vars=2, num_iterations=4, key_base=20)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testWhileWithScopedAllocator(self):
|
||||
group_size = 2
|
||||
group_key = 1
|
||||
@ -284,47 +292,52 @@ class CollectiveOpTest(test.TestCase):
|
||||
del rewrite_options.scoped_allocator_opts.enable_op[:]
|
||||
rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
|
||||
|
||||
with self.session(config=config) as sess:
|
||||
run_ops = []
|
||||
for i in range(group_size):
|
||||
with ops.device('CPU:%d' % i):
|
||||
constant = constant_op.constant(0.)
|
||||
cond = lambda i: math_ops.less(i, 10.)
|
||||
body = lambda i: math_ops.add(i, 1.)
|
||||
input0 = control_flow_ops.while_loop(cond, body, [constant])
|
||||
input1 = math_ops.add(constant, 5)
|
||||
colred0 = collective_ops.all_reduce(input0, group_size, group_key,
|
||||
instance_key0, 'Add', 'Id')
|
||||
colred1 = collective_ops.all_reduce(input1, group_size, group_key,
|
||||
instance_key1, 'Add', 'Id')
|
||||
run_ops.append(math_ops.add_n([colred0, colred1]))
|
||||
results = sess.run(run_ops)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
with self.session(config=config) as sess:
|
||||
run_ops = []
|
||||
for i in range(group_size):
|
||||
with ops.device('CPU:%d' % i):
|
||||
constant = constant_op.constant(0.)
|
||||
cond = lambda i: math_ops.less(i, 10.)
|
||||
body = lambda i: math_ops.add(i, 1.)
|
||||
input0 = control_flow_ops.while_loop(cond, body, [constant])
|
||||
input1 = math_ops.add(constant, 5)
|
||||
colred0 = collective_ops.all_reduce(input0, group_size, group_key,
|
||||
instance_key0, 'Add', 'Id')
|
||||
colred1 = collective_ops.all_reduce(input1, group_size, group_key,
|
||||
instance_key1, 'Add', 'Id')
|
||||
run_ops.append(math_ops.add_n([colred0, colred1]))
|
||||
results = sess.run(run_ops)
|
||||
self.assertEqual(results, [30., 30.])
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveReduceScalar(self):
|
||||
self._testCollectiveReduce(inputs=[0.1, 0.3], expected=0.2,
|
||||
set_graph_key=True)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(inputs=[0.1, 0.3], expected=0.2,
|
||||
set_graph_key=True)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveReduceMaximum(self):
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
|
||||
expected=[10., 20., 30., 40., 50.],
|
||||
set_graph_key=True,
|
||||
instance_key=30,
|
||||
merge_op='Max',
|
||||
final_op='Id')
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
|
||||
expected=[10., 20., 30., 40., 50.],
|
||||
set_graph_key=True,
|
||||
instance_key=30,
|
||||
merge_op='Max',
|
||||
final_op='Id')
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveReduceMinimum(self):
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
|
||||
expected=[1., 2., 3., 4., 5.],
|
||||
set_graph_key=True,
|
||||
instance_key=40,
|
||||
merge_op='Min',
|
||||
final_op='Id')
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveReduce(
|
||||
inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
|
||||
expected=[1., 2., 3., 4., 5.],
|
||||
set_graph_key=True,
|
||||
instance_key=40,
|
||||
merge_op='Min',
|
||||
final_op='Id')
|
||||
|
||||
def _testCollectiveBroadcast(self, in_val):
|
||||
group_key = 1
|
||||
@ -345,13 +358,15 @@ class CollectiveOpTest(test.TestCase):
|
||||
self.assertAllClose(results[0], in_val, rtol=1e-5, atol=1e-5)
|
||||
self.assertAllClose(results[1], in_val, rtol=1e-5, atol=1e-5)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveBroadcast(self):
|
||||
self._testCollectiveBroadcast([0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1])
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveBroadcast([0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1])
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveBroadcastBool(self):
|
||||
self._testCollectiveBroadcast([True, False])
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveBroadcast([True, False])
|
||||
|
||||
def _testCollectiveGather(self, t0, t1, expected, set_graph_key):
|
||||
group_key = 1
|
||||
@ -371,94 +386,101 @@ class CollectiveOpTest(test.TestCase):
|
||||
self.assertAllClose(results[0], expected, rtol=1e-5, atol=1e-5)
|
||||
self.assertAllClose(results[1], expected, rtol=1e-5, atol=1e-5)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveGather(self):
|
||||
self._testCollectiveGather([0, 1, 2, 3, 4, 5, 6, 7],
|
||||
[10, 11, 12, 13, 14, 15, 16, 17],
|
||||
[0, 1, 2, 3, 4, 5, 6, 7,
|
||||
10, 11, 12, 13, 14, 15, 16, 17],
|
||||
True)
|
||||
self._testCollectiveGather([[0, 1, 2, 3], [4, 5, 6, 7]],
|
||||
[[10, 11, 12, 13], [14, 15, 16, 17]],
|
||||
[[0, 1, 2, 3], [4, 5, 6, 7],
|
||||
[10, 11, 12, 13], [14, 15, 16, 17]],
|
||||
True)
|
||||
self._testCollectiveGather([[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
|
||||
[[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
|
||||
[[[0, 1], [2, 3]], [[4, 5], [6, 7]],
|
||||
[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
|
||||
True)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
self._testCollectiveGather([0, 1, 2, 3, 4, 5, 6, 7],
|
||||
[10, 11, 12, 13, 14, 15, 16, 17],
|
||||
[0, 1, 2, 3, 4, 5, 6, 7,
|
||||
10, 11, 12, 13, 14, 15, 16, 17],
|
||||
True)
|
||||
self._testCollectiveGather([[0, 1, 2, 3], [4, 5, 6, 7]],
|
||||
[[10, 11, 12, 13], [14, 15, 16, 17]],
|
||||
[[0, 1, 2, 3], [4, 5, 6, 7],
|
||||
[10, 11, 12, 13], [14, 15, 16, 17]],
|
||||
True)
|
||||
self._testCollectiveGather([[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
|
||||
[[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
|
||||
[[[0, 1], [2, 3]], [[4, 5], [6, 7]],
|
||||
[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
|
||||
True)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveGatherShapeMismatch(self):
|
||||
group_key = 1
|
||||
instance_key = 1
|
||||
t0 = [1, 2, 3, 4]
|
||||
t1 = [5, 6, 7, 8]
|
||||
t2 = [9, 10]
|
||||
with self.session(
|
||||
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = constant_op.constant(t0)
|
||||
c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = constant_op.constant(t1)
|
||||
in2 = constant_op.constant(t2)
|
||||
c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
|
||||
c2 = collective_ops.all_gather(in2, 2, group_key, instance_key)
|
||||
run_options = config_pb2.RunOptions()
|
||||
run_options.experimental.collective_graph_key = 1
|
||||
sess.run([c0, c1], options=run_options)
|
||||
with self.assertRaisesRegexp(errors.InvalidArgumentError,
|
||||
'Shape mismatch'):
|
||||
sess.run([c0, c2], options=run_options)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
with self.session(
|
||||
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = constant_op.constant(t0)
|
||||
c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = constant_op.constant(t1)
|
||||
in2 = constant_op.constant(t2)
|
||||
c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
|
||||
c2 = collective_ops.all_gather(in2, 2, group_key, instance_key)
|
||||
run_options = config_pb2.RunOptions()
|
||||
run_options.experimental.collective_graph_key = 1
|
||||
sess.run([c0, c1], options=run_options)
|
||||
with self.assertRaisesRegexp(errors.InvalidArgumentError,
|
||||
'Shape mismatch'):
|
||||
sess.run([c0, c2], options=run_options)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveGatherShapeMismatchAcrossDevices(self):
|
||||
group_key = 1
|
||||
instance_key = 1
|
||||
t0 = [1, 2, 3, 4]
|
||||
t1 = [5, 6]
|
||||
with self.session(
|
||||
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = constant_op.constant(t0)
|
||||
c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = constant_op.constant(t1)
|
||||
c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
|
||||
run_options = config_pb2.RunOptions()
|
||||
run_options.experimental.collective_graph_key = 1
|
||||
with self.assertRaisesRegexp(errors.InvalidArgumentError,
|
||||
'Shape mismatch'):
|
||||
sess.run([c0, c1], options=run_options)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
with self.session(
|
||||
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = constant_op.constant(t0)
|
||||
c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = constant_op.constant(t1)
|
||||
c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
|
||||
run_options = config_pb2.RunOptions()
|
||||
run_options.experimental.collective_graph_key = 1
|
||||
with self.assertRaisesRegexp(errors.InvalidArgumentError,
|
||||
'Shape mismatch'):
|
||||
sess.run([c0, c1], options=run_options)
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testCollectiveGatherPolymorphicShape(self):
|
||||
t0 = [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
t1 = [10, 11, 12, 13, 14, 15, 16, 17]
|
||||
group_size = 2
|
||||
group_key = 1
|
||||
instance_key = 123
|
||||
with self.session(
|
||||
config=config_pb2.ConfigProto(
|
||||
device_count={'CPU': group_size})) as sess:
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
|
||||
c0 = collective_ops.all_gather(in0, group_size, group_key, instance_key)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
|
||||
c1 = collective_ops.all_gather(in1, group_size, group_key, instance_key)
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
with self.session(
|
||||
config=config_pb2.ConfigProto(
|
||||
device_count={'CPU': group_size})) as sess:
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
|
||||
c0 = collective_ops.all_gather(in0, group_size, group_key,
|
||||
instance_key)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
|
||||
c1 = collective_ops.all_gather(in1, group_size, group_key,
|
||||
instance_key)
|
||||
|
||||
results = sess.run([c0, c1], feed_dict={in0: t0, in1: t1})
|
||||
expected_output = [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 15, 16, 17]
|
||||
self.assertAllClose(results[0], expected_output, rtol=1e-5, atol=1e-5)
|
||||
self.assertAllClose(results[1], expected_output, rtol=1e-5, atol=1e-5)
|
||||
results = sess.run([c0, c1], feed_dict={in0: t0, in1: t1})
|
||||
results_ = sess.run([c0, c1], feed_dict={in0: t0[1:], in1: t1[1:]})
|
||||
|
||||
results_ = sess.run([c0, c1], feed_dict={in0: t0[1:], in1: t1[1:]})
|
||||
expected_output_ = [1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17]
|
||||
self.assertAllClose(results_[0], expected_output_, rtol=1e-5, atol=1e-5)
|
||||
self.assertAllClose(results_[1], expected_output_, rtol=1e-5, atol=1e-5)
|
||||
expected_output = [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 15, 16, 17]
|
||||
self.assertAllClose(results[0], expected_output, rtol=1e-5, atol=1e-5)
|
||||
self.assertAllClose(results[1], expected_output, rtol=1e-5, atol=1e-5)
|
||||
|
||||
expected_output_ = [1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17]
|
||||
self.assertAllClose(results_[0], expected_output_, rtol=1e-5, atol=1e-5)
|
||||
self.assertAllClose(results_[1], expected_output_, rtol=1e-5, atol=1e-5)
|
||||
|
||||
@test_util.run_v2_only
|
||||
def testCollectiveGroupSizeMismatch(self):
|
||||
@ -492,8 +514,17 @@ class CollectiveOpTest(test.TestCase):
|
||||
'but that group has size'):
|
||||
run_all_reduce()
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
@test_util.run_v2_only
|
||||
def testCollectiveTensorsHaveNoDeviceSpecified(self):
|
||||
context._reset_context()
|
||||
cpus = config.list_physical_devices('CPU')
|
||||
self.assertEqual(len(cpus), 1)
|
||||
config.set_logical_device_configuration(cpus[0], [
|
||||
context.LogicalDeviceConfiguration(),
|
||||
context.LogicalDeviceConfiguration()
|
||||
])
|
||||
context.ensure_initialized()
|
||||
|
||||
group_size = 2
|
||||
group_key = 1
|
||||
instance_key = 1
|
||||
@ -517,20 +548,12 @@ class CollectiveOpTest(test.TestCase):
|
||||
|
||||
return results
|
||||
|
||||
with self.session(config=config_pb2.ConfigProto(
|
||||
device_count={'CPU': 2})) as sess:
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = constant_op.constant(1)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = constant_op.constant(3)
|
||||
|
||||
result_op = fn([in0, in1])
|
||||
|
||||
run_options = config_pb2.RunOptions()
|
||||
run_options.experimental.collective_graph_key = 1
|
||||
result = sess.run(result_op, options=run_options)
|
||||
|
||||
self.assertAllClose(result, [2, 2])
|
||||
with ops.device('/CPU:0'):
|
||||
in0 = constant_op.constant(1)
|
||||
with ops.device('/CPU:1'):
|
||||
in1 = constant_op.constant(3)
|
||||
result = fn([in0, in1])
|
||||
self.assertAllClose(result, [2, 2])
|
||||
|
||||
@test_util.run_v2_only
|
||||
def testCollectiveGroupSizeOne(self):
|
||||
@ -548,7 +571,6 @@ class CollectiveOpTest(test.TestCase):
|
||||
in_tensor, group_size, group_key, instance_key)
|
||||
self.assertAllEqual(in_value, gathered_tensor.numpy())
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def testConstantWithScopedAllocator(self):
|
||||
group_size = 2
|
||||
group_key = 1
|
||||
@ -565,21 +587,25 @@ class CollectiveOpTest(test.TestCase):
|
||||
del rewrite_options.scoped_allocator_opts.enable_op[:]
|
||||
rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
|
||||
|
||||
with self.session(config=cfg) as sess:
|
||||
run_ops = []
|
||||
for i in range(group_size):
|
||||
with ops.device('CPU:%d' % i):
|
||||
constant = constant_op.constant(i + 1.)
|
||||
input_tensor1 = array_ops.identity(constant)
|
||||
input_tensor2 = array_ops.identity(constant)
|
||||
reduced_tensor1 = collective_ops.all_reduce(
|
||||
input_tensor1, group_size, group_key, instance_key1, 'Add', 'Id')
|
||||
reduced_tensor2 = collective_ops.all_reduce(
|
||||
input_tensor2, group_size, group_key, instance_key2, 'Add', 'Id')
|
||||
run_ops.append(array_ops.identity(reduced_tensor1))
|
||||
run_ops.append(array_ops.identity(reduced_tensor2))
|
||||
results = sess.run(run_ops)
|
||||
self.assertEqual(results, [3., 3., 3., 3.])
|
||||
# Tests that execute collectives need to be enclosed in graph or tf.function
|
||||
with ops.Graph().as_default():
|
||||
with self.session(config=cfg) as sess:
|
||||
run_ops = []
|
||||
for i in range(group_size):
|
||||
with ops.device('CPU:%d' % i):
|
||||
constant = constant_op.constant(i + 1.)
|
||||
input_tensor1 = array_ops.identity(constant)
|
||||
input_tensor2 = array_ops.identity(constant)
|
||||
reduced_tensor1 = collective_ops.all_reduce(
|
||||
input_tensor1, group_size, group_key, instance_key1, 'Add',
|
||||
'Id')
|
||||
reduced_tensor2 = collective_ops.all_reduce(
|
||||
input_tensor2, group_size, group_key, instance_key2, 'Add',
|
||||
'Id')
|
||||
run_ops.append(array_ops.identity(reduced_tensor1))
|
||||
run_ops.append(array_ops.identity(reduced_tensor2))
|
||||
results = sess.run(run_ops)
|
||||
self.assertEqual(results, [3., 3., 3., 3.])
|
||||
|
||||
@test_util.run_v2_only
|
||||
def testMultipleGroups(self):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user