def test_cardinality_free_data_descriptor_places_data(self):
ds = data_descriptor.CardinalityFreeDataDescriptor(
federated_computation.federated_computation(
lambda x: intrinsics.federated_value(x, placements.CLIENTS),
tf.int32), 1000, computation_types.TensorType(tf.int32))
self.assertEqual(str(ds.type_signature), 'int32@CLIENTS')
@federated_computation.federated_computation(
computation_types.FederatedType(tf.int32,
placements.CLIENTS,
all_equal=True))
def foo(x):
return intrinsics.federated_sum(x)
# Since this DataDescriptor does not specify its cardinality, the number of
# values placed is inferred from the decault setting for the executor.
with executor_test_utils.install_executor(
executor_test_utils.LocalTestExecutorFactory(
default_num_clients=1)):
result = foo(ds)
self.assertEqual(result, 1000)
with executor_test_utils.install_executor(
executor_test_utils.LocalTestExecutorFactory(
default_num_clients=3)):
result = foo(ds)
self.assertEqual(result, 3000)
def test_with_temperature_sensor_example(self, executor):
@computations.tf_computation(computation_types.SequenceType(
tf.float32), tf.float32)
def count_over(ds, t):
return ds.reduce(
np.float32(0),
lambda n, x: n + tf.cast(tf.greater(x, t), tf.float32))
@computations.tf_computation(computation_types.SequenceType(tf.float32)
)
def count_total(ds):
return ds.reduce(np.float32(0.0), lambda n, _: n + 1.0)
@computations.federated_computation(
type_factory.at_clients(computation_types.SequenceType(
tf.float32)), type_factory.at_server(tf.float32))
def comp(temperatures, threshold):
return intrinsics.federated_mean(
intrinsics.federated_map(
count_over,
intrinsics.federated_zip([
temperatures,
intrinsics.federated_broadcast(threshold)
])), intrinsics.federated_map(count_total, temperatures))
with executor_test_utils.install_executor(executor):
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [
tf.data.Dataset.range(10).map(to_float),
tf.data.Dataset.range(20).map(to_float),
tf.data.Dataset.range(30).map(to_float),
]
threshold = 15.0
result = comp(temperatures, threshold)
self.assertAlmostEqual(result, 8.333, places=3)
def test_create_data_descriptor_for_data_backend(self):
class TestDataBackend(data_backend_base.DataBackend):
def __init__(self, value):
self._value = value
async def materialize(self, data, type_spec):
return self._value
data_constant = 1
type_spec = computation_types.TensorType(tf.int32)
def ex_fn(device):
return data_executor.DataExecutor(
eager_tf_executor.EagerTFExecutor(device),
TestDataBackend(data_constant))
factory = executor_test_utils.LocalTestExecutorFactory(
leaf_executor_fn=ex_fn)
@federated_computation.federated_computation(
computation_types.FederatedType(type_spec, placements.CLIENTS))
def foo(dd):
return intrinsics.federated_sum(dd)
with executor_test_utils.install_executor(factory):
uris = [f'foo://bar{i}' for i in range(3)]
dd = data_descriptor.CreateDataDescriptor(uris, type_spec)
result = foo(dd)
self.assertEqual(result, 3)
def test_execution_of_tensorflow(self, executor):
@computations.tf_computation
def comp():
return tf.math.add(5, 5)
with executor_test_utils.install_executor(executor):
result = comp()
self.assertEqual(result, 10)
def test_simple_no_arg_tf_computation_with_int_result(self):
@computations.tf_computation
def comp():
return tf.constant(10)
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
result = comp()
self.assertEqual(result, 10)
def test_one_arg_tf_computation_with_int_param_and_result(self):
@computations.tf_computation(tf.int32)
def comp(x):
return tf.add(x, 10)
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
result = comp(3)
self.assertEqual(result, 13)
def test_three_arg_tf_computation_with_int_params_and_result(self):
@computations.tf_computation(tf.int32, tf.int32, tf.int32)
def comp(x, y, z):
return tf.multiply(tf.add(x, y), z)
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
result = comp(3, 4, 5)
self.assertEqual(result, 35)
def test_execution_of_tensorflow(self):
@tensorflow_computation.tf_computation
def comp():
return tf.math.add(5, 5)
executor = create_test_executor_factory()
with executor_test_utils.install_executor(executor):
result = comp()
self.assertEqual(result, 10)
def test_execution_with_inferred_clients_larger_than_fanout(
self, executor_factory_fn):
@computations.federated_computation(type_factory.at_clients(tf.int32))
def foo(x):
return intrinsics.federated_sum(x)
executor = executor_factory_fn(max_fanout=3)
with executor_test_utils.install_executor(executor):
result = foo([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
self.assertEqual(result, 55)
def test_tf_computation_with_dataset_params_and_int_result(self):
@computations.tf_computation(computation_types.SequenceType(tf.int32))
def comp(ds):
return ds.reduce(np.int32(0), lambda x, y: x + y)
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
ds = tf.data.Dataset.range(10).map(lambda x: tf.cast(x, tf.int32))
result = comp(ds)
self.assertEqual(result, 45)
def test_tuple_argument_can_accept_unnamed_elements(self):
@computations.tf_computation(tf.int32, tf.int32)
def foo(x, y):
return x + y
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
# pylint:disable=no-value-for-parameter
result = foo(structure.Struct([(None, 2), (None, 3)]))
# pylint:enable=no-value-for-parameter
self.assertEqual(result, 5)
def test_execution_of_temperature_sensor_example(self, executor):
comp = _temperature_sensor_example_next_fn()
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [
tf.data.Dataset.range(10).map(to_float),
tf.data.Dataset.range(20).map(to_float),
tf.data.Dataset.range(30).map(to_float),
]
threshold = 15.0
with executor_test_utils.install_executor(executor):
result = comp(temperatures, threshold)
self.assertAlmostEqual(result, 8.333, places=3)
def test_local_executor_multi_gpus_iter_dataset(self, tf_device):
tf_devices = tf.config.list_logical_devices(tf_device)
server_tf_device = None if not tf_devices else tf_devices[0]
gpu_devices = tf.config.list_logical_devices('GPU')
local_executor = executor_stacks.local_executor_factory(
server_tf_device=server_tf_device, client_tf_devices=gpu_devices)
with executor_test_utils.install_executor(local_executor):
parallel_client_run = _create_tff_parallel_clients_with_iter_dataset()
client_data = [
tf.data.Dataset.range(10),
tf.data.Dataset.range(10).map(lambda x: x + 1)
]
client_results = parallel_client_run(client_data)
self.assertEqual(client_results, [np.int64(46), np.int64(56)])
def test_tf_computation_with_structured_result(self):
@computations.tf_computation
def comp():
return collections.OrderedDict([
('a', tf.constant(10)),
('b', tf.constant(20)),
])
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
result = comp()
self.assertIsInstance(result, collections.OrderedDict)
self.assertDictEqual(result, {'a': 10, 'b': 20})
def test_non_federated(self):
ds = data_descriptor.DataDescriptor(
tensorflow_computation.tf_computation(
lambda x: tf.cast(x + 10, tf.float32), tf.int32), 20,
computation_types.TensorType(tf.int32))
self.assertEqual(str(ds.type_signature), 'float32')
@tensorflow_computation.tf_computation(tf.float32)
def foo(x):
return x * 20.0
with executor_test_utils.install_executor(
executor_test_utils.LocalTestExecutorFactory()):
result = foo(ds)
self.assertEqual(result, 600.0)
def test_comp_none(self):
ds = data_descriptor.DataDescriptor(
None, [1, 2, 3],
computation_types.FederatedType(tf.int32, placements.CLIENTS), 3)
self.assertEqual(str(ds.type_signature), '{int32}@CLIENTS')
@federated_computation.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
return intrinsics.federated_sum(x)
with executor_test_utils.install_executor(
executor_test_utils.LocalTestExecutorFactory()):
result = foo(ds)
self.assertEqual(result, 6)
def test_changing_cardinalities_across_calls(self):
@computations.federated_computation(type_factory.at_clients(tf.int32))
def comp(x):
return x
five_ints = list(range(5))
ten_ints = list(range(10))
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
five = comp(five_ints)
ten = comp(ten_ints)
self.assertEqual(five, five_ints)
self.assertEqual(ten, ten_ints)
def test_conflicting_cardinalities_within_call(self):
@computations.federated_computation([
computation_types.at_clients(tf.int32),
computation_types.at_clients(tf.int32),
])
def comp(x):
return x
five_ints = list(range(5))
ten_ints = list(range(10))
executor = executor_stacks.local_executor_factory()
with executor_test_utils.install_executor(executor):
with self.assertRaisesRegex(ValueError,
'Conflicting cardinalities'):
comp([five_ints, ten_ints])
def test_federated(self):
ds = data_descriptor.DataDescriptor(
federated_computation.federated_computation(
lambda x: intrinsics.federated_value(x, placements.CLIENTS),
tf.int32), 1000, computation_types.TensorType(tf.int32), 3)
self.assertEqual(str(ds.type_signature), 'int32@CLIENTS')
@federated_computation.federated_computation(
computation_types.FederatedType(tf.int32,
placements.CLIENTS,
all_equal=True))
def foo(x):
return intrinsics.federated_sum(x)
with executor_test_utils.install_executor(
executor_test_utils.LocalTestExecutorFactory()):
result = foo(ds)
self.assertEqual(result, 3000)
def test_local_executor_multi_gpus_dataset_reduce(self, tf_device):
tf_devices = tf.config.list_logical_devices(tf_device)
server_tf_device = None if not tf_devices else tf_devices[0]
gpu_devices = tf.config.list_logical_devices('GPU')
local_executor = executor_stacks.local_executor_factory(
server_tf_device=server_tf_device, client_tf_devices=gpu_devices)
with executor_test_utils.install_executor(local_executor):
parallel_client_run = _create_tff_parallel_clients_with_dataset_reduce()
client_data = [
tf.data.Dataset.range(10),
tf.data.Dataset.range(10).map(lambda x: x + 1)
]
# TODO(b/159180073): merge this one into iter dataset test when the
# dataset reduce function can be correctly used for GPU device.
with self.assertRaisesRegex(
ValueError,
'Detected dataset reduce op in multi-GPU TFF simulation.*'):
parallel_client_run(client_data)
