def test_with_federated_map_and_broadcast(self): eager_ex = eager_executor.EagerExecutor() federated_ex = federated_executor.FederatedExecutor({ None: eager_ex, placement_literals.SERVER: eager_ex, placement_literals.CLIENTS: [eager_ex for _ in range(3)] }) ex = lambda_executor.LambdaExecutor(federated_ex) loop = asyncio.get_event_loop() @computations.tf_computation(tf.int32) def add_one(x): return x + 1 @computations.federated_computation( type_constructors.at_server(tf.int32)) def comp(x): return intrinsics.federated_map(add_one, intrinsics.federated_broadcast(x)) v1 = loop.run_until_complete(ex.create_value(comp)) v2 = loop.run_until_complete( ex.create_value(10, type_constructors.at_server(tf.int32))) v3 = loop.run_until_complete(ex.create_call(v1, v2)) result = loop.run_until_complete(v3.compute()) self.assertCountEqual([x.numpy() for x in result], [11, 11, 11])
def test_with_federated_apply(self): eager_ex = eager_executor.EagerExecutor() federated_ex = federated_executor.FederatedExecutor({ None: eager_ex, placement_literals.SERVER: eager_ex }) ex = lambda_executor.LambdaExecutor(federated_ex) loop = asyncio.get_event_loop() @computations.tf_computation(tf.int32) def add_one(x): return x + 1 @computations.federated_computation( type_constructors.at_server(tf.int32)) def comp(x): return intrinsics.federated_apply(add_one, x) v1 = loop.run_until_complete(ex.create_value(comp)) v2 = loop.run_until_complete( ex.create_value(10, type_constructors.at_server(tf.int32))) v3 = loop.run_until_complete(ex.create_call(v1, v2)) result = loop.run_until_complete(v3.compute()) self.assertEqual(result.numpy(), 11)
def test_executor_create_value_with_valid_intrinsic_def(self): loop = asyncio.get_event_loop() ex = _make_test_executor() val = loop.run_until_complete( ex.create_value( intrinsic_defs.FEDERATED_APPLY, computation_types.FunctionType([ type_constructors.unary_op(tf.int32), type_constructors.at_server(tf.int32) ], type_constructors.at_server(tf.int32)))) self.assertIsInstance(val, federated_executor.FederatedExecutorValue) self.assertEqual(str(val.type_signature), '(<(int32 -> int32),int32@SERVER> -> int32@SERVER)') self.assertIs(val.internal_representation, intrinsic_defs.FEDERATED_APPLY)
def test_with_removal_of_identity_mapping(self): @computations.federated_computation( type_constructors.at_server(tf.int32)) def comp(x): return intrinsics.federated_apply(_identity, x) def transformation_fn(x): x, _ = transformations.remove_mapped_or_applied_identity(x) return x self.assertEqual(_test_create_value(comp, transformation_fn), '(FEDERATED_arg -> FEDERATED_arg)')
def test_executor_create_value_with_server_int(self): loop = asyncio.get_event_loop() ex = _make_test_executor() val = loop.run_until_complete( ex.create_value(10, type_constructors.at_server(tf.int32))) self.assertIsInstance(val, federated_executor.FederatedExecutorValue) self.assertEqual(str(val.type_signature), 'int32@SERVER') self.assertIsInstance(val.internal_representation, list) self.assertLen(val.internal_representation, 1) self.assertIsInstance(val.internal_representation[0], eager_executor.EagerValue) self.assertEqual( val.internal_representation[0].internal_representation.numpy(), 10)
def test_with_inlining_of_blocks(self): @computations.federated_computation(type_constructors.at_server(tf.int32)) def comp(x): return intrinsics.federated_zip([x, x]) # TODO(b/134543154): Slide in something more powerful so that this test # doesn't break when the implementation changes; for now, this will do. def transformation_fn(x): x, _ = transformations.remove_mapped_or_applied_identity(x) x, _ = transformations.inline_block_locals(x) x, _ = transformations.replace_selection_from_tuple_with_element(x) return x self.assertIn('federated_zip_at_server(<FEDERATED_arg,FEDERATED_arg>)', _test_create_value(comp, transformation_fn))
def test_with_temperature_sensor_example(self): @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_constructors.at_clients( computation_types.SequenceType(tf.float32)), type_constructors.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)) num_clients = 3 set_default_executor.set_default_executor( executor_stacks.create_local_executor(num_clients)) 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) set_default_executor.set_default_executor()
def test_with_incomplete_temperature_sensor_example(self): @computations.federated_computation( type_constructors.at_clients( computation_types.SequenceType(tf.float32)), type_constructors.at_server(tf.float32)) def comp(temperatures, threshold): @computations.tf_computation( computation_types.SequenceType(tf.float32), tf.float32) def count(ds, t): return ds.reduce( np.int32(0), lambda n, x: n + tf.cast(tf.greater(x, t), tf.int32)) return intrinsics.federated_map( count, intrinsics.federated_zip( [temperatures, intrinsics.federated_broadcast(threshold)])) num_clients = 10 set_default_executor.set_default_executor( executor_stacks.create_local_executor(num_clients)) temperatures = [ tf.data.Dataset.range(1000).map(lambda x: tf.cast(x, tf.float32)) for _ in range(num_clients) ] threshold = 100.0 result = comp(temperatures, threshold) self.assertCountEqual([x.numpy() for x in result], [899 for _ in range(num_clients)]) set_default_executor.set_default_executor()
def test_federated_mean_with_floats(self): loop = asyncio.get_event_loop() ex = _make_test_executor(4) v1 = loop.run_until_complete( ex.create_value([1.0, 2.0, 3.0, 4.0], type_constructors.at_clients(tf.float32))) self.assertEqual(str(v1.type_signature), '{float32}@CLIENTS') v2 = loop.run_until_complete( ex.create_value( intrinsic_defs.FEDERATED_MEAN, computation_types.FunctionType( type_constructors.at_clients(tf.float32), type_constructors.at_server(tf.float32)))) self.assertEqual(str(v2.type_signature), '({float32}@CLIENTS -> float32@SERVER)') v3 = loop.run_until_complete(ex.create_call(v2, v1)) self.assertEqual(str(v3.type_signature), 'float32@SERVER') result = loop.run_until_complete(v3.compute()) self.assertEqual(result.numpy(), 2.5)
def test_federated_weighted_mean_with_floats(self): loop = asyncio.get_event_loop() ex = _make_test_executor(num_clients=4, use_lambda_executor=True) v1 = loop.run_until_complete( ex.create_value([1.0, 2.0, 3.0, 4.0], type_constructors.at_clients(tf.float32))) self.assertEqual(str(v1.type_signature), '{float32}@CLIENTS') v2 = loop.run_until_complete( ex.create_value([5.0, 10.0, 3.0, 2.0], type_constructors.at_clients(tf.float32))) self.assertEqual(str(v2.type_signature), '{float32}@CLIENTS') v3 = loop.run_until_complete( ex.create_tuple( anonymous_tuple.AnonymousTuple([(None, v1), (None, v2)]))) self.assertEqual(str(v3.type_signature), '<{float32}@CLIENTS,{float32}@CLIENTS>') v4 = loop.run_until_complete( ex.create_value( intrinsic_defs.FEDERATED_WEIGHTED_MEAN, computation_types.FunctionType([ type_constructors.at_clients(tf.float32), type_constructors.at_clients(tf.float32) ], type_constructors.at_server(tf.float32)))) self.assertEqual( str(v4.type_signature), '(<{float32}@CLIENTS,{float32}@CLIENTS> -> float32@SERVER)') v5 = loop.run_until_complete(ex.create_call(v4, v3)) self.assertEqual(str(v5.type_signature), 'float32@SERVER') result = loop.run_until_complete(v5.compute()) self.assertAlmostEqual(result.numpy(), 2.1, places=3)
# # Actual implementations might vary. # # Type signature: <{T}@CLIENTS,U,(<U,T>->U),(<U,U>->U),(U->R)> -> R@SERVER FEDERATED_AGGREGATE = IntrinsicDef( 'FEDERATED_AGGREGATE', 'federated_aggregate', computation_types.FunctionType(parameter=[ type_constructors.at_clients(computation_types.AbstractType('T')), computation_types.AbstractType('U'), type_constructors.reduction_op(computation_types.AbstractType('U'), computation_types.AbstractType('T')), type_constructors.binary_op(computation_types.AbstractType('U')), computation_types.FunctionType(computation_types.AbstractType('U'), computation_types.AbstractType('R')) ], result=type_constructors.at_server( computation_types.AbstractType('R')))) # Applies a given function to a value on the server. # # Type signature: <(T->U),T@SERVER> -> U@SERVER FEDERATED_APPLY = IntrinsicDef( 'FEDERATED_APPLY', 'federated_apply', computation_types.FunctionType(parameter=[ computation_types.FunctionType(computation_types.AbstractType('T'), computation_types.AbstractType('U')), type_constructors.at_server(computation_types.AbstractType('T')), ], result=type_constructors.at_server( computation_types.AbstractType('U')))) # Broadcasts a server item to all clients.
def test_at_server(self): self.assertEqual(str(type_constructors.at_server(tf.bool)), 'bool@SERVER')