def test_serialize_tensorflow_with_data_set_sum_lambda(self):
def _legacy_dataset_reducer_example(ds):
return ds.reduce(np.int64(0), lambda x, y: x + y)
comp = tensorflow_serialization.serialize_py_fn_as_tf_computation(
_legacy_dataset_reducer_example,
computation_types.SequenceType(tf.int64),
context_stack_impl.context_stack)
self.assertEqual(str(type_serialization.deserialize_type(comp.type)),
'(int64* -> int64)')
self.assertEqual(comp.WhichOneof('computation'), 'tensorflow')
parameter = tf.data.Dataset.range(5)
results = tf.Session().run(
tf.import_graph_def(
serialization_utils.unpack_graph_def(
comp.tensorflow.graph_def),
{
comp.tensorflow.parameter.sequence.iterator_string_handle_name:
(parameter.make_one_shot_iterator().string_handle())
}, [comp.tensorflow.result.tensor.tensor_name]))
self.assertEqual(results, [10])
def test_federated_collect_with_map_call(self):
@computations.tf_computation()
def make_dataset():
return tf.data.Dataset.range(5)
@computations.tf_computation(computation_types.SequenceType(tf.int64))
def foo(x):
return x.reduce(tf.constant(0, dtype=tf.int64), lambda a, b: a + b)
@computations.federated_computation()
def bar():
x = intrinsics.federated_value(make_dataset(),
placement_literals.CLIENTS)
return intrinsics.federated_map(
foo,
intrinsics.federated_collect(intrinsics.federated_map(foo, x)))
result = _run_test_comp_produces_federated_value(self,
bar,
num_clients=5)
self.assertEqual(result.numpy(), 50)
def test_assemble_result_from_graph_with_sequence_of_odicts(self):
type_spec = computation_types.SequenceType(
collections.OrderedDict([('X', tf.int32), ('Y', tf.int32)]))
binding = pb.TensorFlow.Binding(sequence=pb.TensorFlow.SequenceBinding(
iterator_string_handle_name='foo'))
data_set = tf.data.Dataset.from_tensors({
'X': tf.constant(1),
'Y': tf.constant(2)
})
it = data_set.make_one_shot_iterator()
output_map = {'foo': it.string_handle()}
result = graph_utils.assemble_result_from_graph(
type_spec, binding, output_map)
self.assertIsInstance(result, graph_utils.DATASET_REPRESENTATION_TYPES)
self.assertEqual(
str(result.output_types),
'OrderedDict([(\'X\', tf.int32), (\'Y\', tf.int32)])')
self.assertEqual(
str(result.output_shapes),
'OrderedDict([(\'X\', TensorShape([])), (\'Y\', TensorShape([]))])'
)
def _create_called_sequence_map(fn, arg):
r"""Creates a to call a sequence map.
Call
/ \
Intrinsic Tuple
/ \
Comp Comp
Args:
fn: A functional `computation_building_blocks.ComputationBuildingBlock` to
use as the function.
arg: A `computation_building_blocks.ComputationBuildingBlock` to use as the
argument.
Returns:
A `computation_building_blocks.Call`.
Raises:
TypeError: If `fn` or `arg` is not a
`computation_building_blocks.ComputationBuildingBlock` or if `fn` has a
parameter type that is not assignable from `arg` type.
"""
py_typecheck.check_type(
fn, computation_building_blocks.ComputationBuildingBlock)
py_typecheck.check_type(
arg, computation_building_blocks.ComputationBuildingBlock)
if not type_utils.is_assignable_from(fn.parameter_type,
arg.type_signature.element):
raise TypeError(
'The parameter of the function is of type {}, and the argument is of '
'an incompatible type {}.'.format(str(fn.parameter_type),
str(arg.type_signature.element)))
result_type = computation_types.SequenceType(fn.type_signature.result)
intrinsic_type = computation_types.FunctionType(
[fn.type_signature, arg.type_signature], result_type)
intrinsic = computation_building_blocks.Intrinsic(
intrinsic_defs.SEQUENCE_MAP.uri, intrinsic_type)
tup = computation_building_blocks.Tuple((fn, arg))
return computation_building_blocks.Call(intrinsic, tup)
def test_infer_cardinalities_success_structure(self):
foo = cardinalities_utils.infer_cardinalities(
structure.Struct([('A', [1, 2, 3]),
('B',
structure.Struct([('C', [[1, 2], [3, 4], [5,
6]]),
('D', [True, False,
True])]))]),
computation_types.StructType([
('A',
computation_types.FederatedType(tf.int32,
placements.CLIENTS)),
('B',
[('C',
computation_types.FederatedType(
computation_types.SequenceType(tf.int32),
placements.CLIENTS)),
('D',
computation_types.FederatedType(tf.bool,
placements.CLIENTS))])
]))
self.assertDictEqual(foo, {placements.CLIENTS: 3})
def test_graph_mode_dataset_fails_well(self):
sequence_type = computation_types.SequenceType(tf.int32)
federated_type = computation_types.FederatedType(
sequence_type, placements.CLIENTS)
with tf.Graph().as_default():
@computations.tf_computation(sequence_type)
def foo(z):
value1 = z.reduce(0, lambda x, y: x + y)
return value1
@computations.federated_computation(federated_type)
def bar(x):
return intrinsics.federated_map(foo, x)
ds1 = tf.data.Dataset.from_tensor_slices([10, 20])
ds2 = tf.data.Dataset.from_tensor_slices([30, 40])
with self.assertRaisesRegexp(
ValueError,
'outside of eager mode is not currently supported.'):
bar([ds1, ds2])
class IsStructureOfIntegersTest(parameterized.TestCase):
@parameterized.named_parameters(
('int', tf.int32),
('ints', ([tf.int32, tf.int32], )),
('federated_int_at_clients',
computation_types.FederatedType(tf.int32,
placement_literals.CLIENTS)),
)
def test_returns_true(self, type_spec):
self.assertTrue(type_analysis.is_structure_of_integers(type_spec))
@parameterized.named_parameters(
('bool', tf.bool),
('string', tf.string),
('int_and_bool', ([tf.int32, tf.bool], )),
('sequence_of_ints', computation_types.SequenceType(tf.int32)),
('placement', computation_types.PlacementType()),
('function', computation_types.FunctionType(tf.int32, tf.int32)),
('abstract', computation_types.AbstractType('T')),
)
def test_returns_false(self, type_spec):
self.assertFalse(type_analysis.is_structure_of_integers(type_spec))
def test_infer_cardinalities_success_anonymous_tuple(self):
foo = cardinalities_utils.infer_cardinalities(
anonymous_tuple.AnonymousTuple([
('A', [1, 2, 3]),
('B',
anonymous_tuple.AnonymousTuple([('C', [[1, 2], [3, 4], [5,
6]]),
('D', [True, False, True])]))
]),
computation_types.to_type([
('A',
computation_types.FederatedType(tf.int32,
placement_literals.CLIENTS)),
('B', [('C',
computation_types.FederatedType(
computation_types.SequenceType(tf.int32),
placement_literals.CLIENTS)),
('D',
computation_types.FederatedType(
tf.bool, placement_literals.CLIENTS))])
]))
self.assertDictEqual(foo, {placement_literals.CLIENTS: 3})
def test_serialize_tensorflow_with_data_set_sum_lambda(self):
def _legacy_dataset_reducer_example(ds):
return ds.reduce(np.int64(0), lambda x, y: x + y)
comp, extra_type_spec = tensorflow_serialization.serialize_py_fn_as_tf_computation(
_legacy_dataset_reducer_example,
computation_types.SequenceType(tf.int64),
context_stack_impl.context_stack)
self.assertEqual(
str(type_serialization.deserialize_type(comp.type)),
'(int64* -> int64)')
self.assertEqual(str(extra_type_spec), '(int64* -> int64)')
self.assertEqual(comp.WhichOneof('computation'), 'tensorflow')
parameter = tf.data.Dataset.range(5)
results = tf.compat.v1.Session().run(
tf.import_graph_def(
serialization_utils.unpack_graph_def(comp.tensorflow.graph_def), {
comp.tensorflow.parameter.sequence.variant_tensor_name:
tf.data.experimental.to_variant(parameter)
}, [comp.tensorflow.result.tensor.tensor_name]))
self.assertEqual(results, [10])
def test_computation_with_int_sequence_raises(self):
ds1_shape = tf.TensorShape([None])
sequence_type = computation_types.SequenceType(
computation_types.TensorType(tf.int32, ds1_shape))
federated_type = computation_types.FederatedType(sequence_type,
placements.CLIENTS)
@computations.tf_computation(sequence_type)
def foo(z):
value1 = z.reduce(0, lambda x, y: x + y)
return value1
@computations.federated_computation(federated_type)
def bar(x):
return intrinsics.federated_map(foo, x)
ds1 = tf.data.Dataset.from_tensor_slices([10, 20]).batch(1)
ds2 = tf.data.Dataset.from_tensor_slices([30, 40]).batch(1)
with self.assertRaisesRegexp(ValueError, 'Please pass a list'):
bar(ds1)
bar(ds2)
def _serialize_sequence_value(
value: Union[type_conversions.TF_DATASET_REPRESENTATION_TYPES],
type_spec: computation_types.SequenceType) -> _SerializeReturnType:
"""Serializes a `tf.data.Dataset` value into `executor_pb2.Value`.
Args:
value: A `tf.data.Dataset`, or equivalent.
type_spec: A `computation_types.Type` specifying the TFF sequence type of
`value.`
Returns:
A tuple `(value_proto, type_spec)` in which `value_proto` is an instance
of `executor_pb2.Value` with the serialized content of `value`, and
`type_spec` is the type of the serialized value.
"""
if not isinstance(value, type_conversions.TF_DATASET_REPRESENTATION_TYPES):
raise TypeError(
'Cannot serialize Python type {!s} as TFF type {!s}.'.format(
py_typecheck.type_string(type(value)),
type_spec if type_spec is not None else 'unknown'))
value_type = computation_types.SequenceType(
computation_types.to_type(value.element_spec))
if not type_spec.is_assignable_from(value_type):
raise TypeError(
'Cannot serialize dataset with elements of type {!s} as TFF type {!s}.'
.format(value_type, type_spec if type_spec is not None else 'unknown'))
# TFF must store the type spec here because TF will lose the ordering of the
# names for `tf.data.Dataset` that return elements of
# `collections.abc.Mapping` type. This allows TFF to preserve and restore the
# key ordering upon deserialization.
element_type = computation_types.to_type(value.element_spec)
return executor_pb2.Value(
sequence=executor_pb2.Value.Sequence(
zipped_saved_model=_serialize_dataset(value),
element_type=type_serialization.serialize_type(
element_type))), type_spec
class CreateIdentityTest(parameterized.TestCase):
# pyformat: disable
@parameterized.named_parameters(
('int', computation_types.TensorType(tf.int32), 10),
('unnamed_tuple', computation_types.StructType([tf.int32, tf.float32]),
structure.Struct([(None, 10), (None, 10.0)])),
('named_tuple',
computation_types.StructType([
('a', tf.int32), ('b', tf.float32)
]), structure.Struct([('a', 10), ('b', 10.0)])),
('sequence', computation_types.SequenceType(tf.int32), [10] * 3),
)
# pyformat: enable
def test_returns_computation(self, type_signature, value):
proto, _ = tensorflow_computation_factory.create_identity(
type_signature)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
expected_type = type_factory.unary_op(type_signature)
self.assertEqual(actual_type, expected_type)
actual_result = test_utils.run_tensorflow(proto, value)
self.assertEqual(actual_result, value)
@parameterized.named_parameters(
('none', None),
('federated_type', computation_types.at_server(tf.int32)),
)
def test_raises_type_error(self, type_signature):
with self.assertRaises(TypeError):
tensorflow_computation_factory.create_identity(type_signature)
def test_feeds_and_fetches_different(self):
proto, _ = tensorflow_computation_factory.create_identity(
computation_types.TensorType(tf.int32))
self.assertNotEqual(proto.tensorflow.parameter,
proto.tensorflow.result)
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_sequence_map_tf_dataset(self):
ds = tf.data.Dataset.range(3)
map_fn = computations.tf_computation(lambda x: x + 2, tf.int64)
sequence_map_val = _make_sequence_map_value(self._sequence_executor,
tf.int64, tf.int64)
ds_val = _run_sync(
self._sequence_executor.create_value(
ds, computation_types.SequenceType(tf.int64)))
map_fn_val = _run_sync(
self._sequence_executor.create_value(
computation_impl.ComputationImpl.get_proto(map_fn),
map_fn.type_signature))
arg_val = _run_sync(
self._sequence_executor.create_struct([map_fn_val, ds_val]))
result_val = _run_sync(
self._sequence_executor.create_call(sequence_map_val, arg_val))
self.assertIsInstance(result_val,
sequence_executor.SequenceExecutorValue)
self.assertEqual(str(result_val.type_signature), 'int64*')
self.assertIsInstance(result_val.internal_representation,
sequence_executor._SequenceFromMap)
result = list(_run_sync(result_val.compute()))
self.assertListEqual(result, [2, 3, 4])
def test_serialize_deserialize_sequence_of_namedtuples(self):
test_tuple_type = collections.namedtuple('TestTuple', ['a', 'b', 'c'])
def make_test_tuple(x):
return test_tuple_type(
a=x * 2, b=tf.cast(x, tf.int32), c=tf.cast(x - 1, tf.float32))
ds = tf.data.Dataset.range(5).map(make_test_tuple)
element_type = computation_types.to_type(
test_tuple_type(tf.int64, tf.int32, tf.float32))
sequence_type = computation_types.SequenceType(element=element_type)
value_proto, value_type = executor_serialization.serialize_value(
ds, sequence_type)
self.assertIsInstance(value_proto, executor_pb2.Value)
self.assertEqual(value_type, sequence_type)
y, type_spec = executor_serialization.deserialize_value(value_proto)
self.assert_types_equivalent(type_spec, sequence_type)
actual_values = list(y)
expected_values = [
test_tuple_type(a=x * 2, b=x, c=x - 1.) for x in range(5)
]
for actual, expected in zip(actual_values, expected_values):
self.assertAllClose(actual, expected)
def test_get_size_info(self, num_clients):
@computations.federated_computation(
type_factory.at_clients(computation_types.SequenceType(tf.float32)),
type_factory.at_server(tf.float32))
def comp(temperatures, threshold):
client_data = [temperatures, intrinsics.federated_broadcast(threshold)]
result_map = intrinsics.federated_map(
count_over, intrinsics.federated_zip(client_data))
count_map = intrinsics.federated_map(count_total, temperatures)
return intrinsics.federated_mean(result_map, count_map)
factory = executor_stacks.sizing_executor_factory(num_clients=num_clients)
default_executor.set_default_executor(factory)
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [tf.data.Dataset.range(10).map(to_float)] * num_clients
threshold = 15.0
comp(temperatures, threshold)
# Each client receives a tf.float32 and uploads two tf.float32 values.
expected_broadcast_bits = [num_clients * 32]
expected_aggregate_bits = [num_clients * 32 * 2]
expected_broadcast_history = {
(('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients
}
expected_aggregate_history = {
(('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients * 2
}
size_info = factory.get_size_info()
self.assertEqual(expected_broadcast_history, size_info.broadcast_history)
self.assertEqual(expected_aggregate_history, size_info.aggregate_history)
self.assertEqual(expected_broadcast_bits, size_info.broadcast_bits)
self.assertEqual(expected_aggregate_bits, size_info.aggregate_bits)
def test_batching_namedtuple_dataset(self):
batch_type = collections.namedtuple('Batch', ['x', 'y'])
federated_sequence_type = computation_types.FederatedType(
computation_types.SequenceType(
batch_type(
x=computation_types.TensorType(tf.float32, [None, 2]),
y=computation_types.TensorType(tf.float32, [None, 1]))),
placements.CLIENTS,
all_equal=False)
@computations.tf_computation(federated_sequence_type.member)
def test_batch_select_and_reduce(z):
i = z.map(lambda x: x.y)
return i.reduce(0., lambda x, y: x + tf.reduce_sum(y))
@computations.federated_computation(federated_sequence_type)
def map_y_sum(x):
return intrinsics.federated_map(test_batch_select_and_reduce, x)
ds = tf.data.Dataset.from_tensor_slices({
'x': [[1., 2.], [3., 4.]],
'y': [[5.], [6.]]
}).batch(1)
self.assertEqual(map_y_sum([ds] * 5), [np.array([[11.]])] * 5)
class IsSumCompatibleTest(parameterized.TestCase):
@parameterized.named_parameters([
('tensor_type', computation_types.TensorType(tf.int32)),
('tuple_type_int', computation_types.StructType([tf.int32, tf.int32],)),
('tuple_type_float',
computation_types.StructType([tf.complex128, tf.float32, tf.float64])),
('federated_type',
computation_types.FederatedType(tf.int32, placement_literals.CLIENTS)),
])
def test_positive_examples(self, type_spec):
self.assertTrue(type_analysis.is_sum_compatible(type_spec))
@parameterized.named_parameters([
('tensor_type_bool', computation_types.TensorType(tf.bool)),
('tensor_type_string', computation_types.TensorType(tf.string)),
('tuple_type', computation_types.StructType([tf.int32, tf.bool])),
('sequence_type', computation_types.SequenceType(tf.int32)),
('placement_type', computation_types.PlacementType()),
('function_type', computation_types.FunctionType(tf.int32, tf.int32)),
('abstract_type', computation_types.AbstractType('T')),
])
def test_negative_examples(self, type_spec):
self.assertFalse(type_analysis.is_sum_compatible(type_spec))
def test_dataset_computation_where_client_data_is_ordered_dicts(self):
client_data = from_tensor_slices_client_data.TestClientData(
TEST_DATA_WITH_ORDEREDDICTS)
dataset_computation = client_data.dataset_computation
self.assertIsInstance(dataset_computation, computation_base.Computation)
expected_dataset_comp_type_signature = computation_types.FunctionType(
computation_types.to_type(tf.string),
computation_types.SequenceType(
collections.OrderedDict([
('x',
computation_types.TensorType(
client_data.element_type_structure['x'].dtype,
tf.TensorShape(2))),
('y',
computation_types.TensorType(
client_data.element_type_structure['y'].dtype, None)),
('z',
computation_types.TensorType(
client_data.element_type_structure['z'].dtype, None))
])))
self.assertTrue(
dataset_computation.type_signature.is_equivalent_to(
expected_dataset_comp_type_signature))
# Iterate over each client, invoking the computation and ensuring
# we received a tf.data.Dataset with the correct data.
for client_id in TEST_DATA_WITH_ORDEREDDICTS:
dataset = dataset_computation(client_id)
self.assertIsInstance(dataset, tf.data.Dataset)
expected_dataset = tf.data.Dataset.from_tensor_slices(
TEST_DATA_WITH_ORDEREDDICTS[client_id])
self.assertSameDatasetsOfDicts(expected_dataset, dataset)
class CreateReplicateInputTest(parameterized.TestCase):
@parameterized.named_parameters(
('int', computation_types.TensorType(tf.int32), 3, 10),
('float', computation_types.TensorType(tf.float32), 3, 10.0),
('unnamed_tuple',
computation_types.NamedTupleType([tf.int32, tf.float32]), 3,
anonymous_tuple.AnonymousTuple([(None, 10), (None, 10.0)])),
('named_tuple',
computation_types.NamedTupleType([
('a', tf.int32), ('b', tf.float32)
]), 3, anonymous_tuple.AnonymousTuple([('a', 10), ('b', 10.0)])),
('sequence', computation_types.SequenceType(tf.int32), 3, [10] * 3),
)
def test_returns_computation(self, type_signature, count, value):
proto = tensorflow_computation_factory.create_replicate_input(
type_signature, count)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
expected_type = computation_types.FunctionType(
type_signature, [type_signature] * count)
self.assertEqual(actual_type, expected_type)
actual_result = test_utils.run_tensorflow(proto, value)
expected_result = anonymous_tuple.AnonymousTuple([(None, value)] *
count)
self.assertEqual(actual_result, expected_result)
@parameterized.named_parameters(
('none_type', None, 3),
('none_count', computation_types.TensorType(tf.int32), None),
('federated_type', type_factory.at_server(tf.int32), 3),
)
def test_raises_type_error(self, type_signature, count):
with self.assertRaises(TypeError):
tensorflow_computation_factory.create_replicate_input(
type_signature, count)
def test_call_returned_directly_creates_canonical_form(self):
@tff.federated_computation
def init_fn():
return tff.federated_value(42, tff.SERVER)
@tff.federated_computation(tff.FederatedType(tf.int32, tff.SERVER),
tff.FederatedType(
tff.SequenceType(tf.float32),
tff.CLIENTS))
def next_fn(server_state, client_data):
broadcast_state = tff.federated_broadcast(server_state)
@tff.tf_computation(tf.int32, tff.SequenceType(tf.float32))
@tf.function
def some_transform(x, y):
del y # Unused
return x + 1
client_update = tff.federated_map(some_transform,
(broadcast_state, client_data))
aggregate_update = tff.federated_sum(client_update)
server_output = tff.federated_value(1234, tff.SERVER)
return aggregate_update, server_output
@tff.federated_computation(
tff.FederatedType(tf.int32, tff.SERVER),
tff.FederatedType(computation_types.SequenceType(tf.float32),
tff.CLIENTS))
def nested_next_fn(server_state, client_data):
return next_fn(server_state, client_data)
iterative_process = computation_utils.IterativeProcess(
init_fn, nested_next_fn)
cf = canonical_form_utils.get_canonical_form_for_iterative_process(
iterative_process)
self.assertIsInstance(cf, canonical_form.CanonicalForm)
def test_stamp_parameter_in_graph_with_bool_sequence(self):
with tf.Graph().as_default():
x = self._checked_stamp_parameter('foo',
computation_types.SequenceType(tf.bool))
self.assertIsInstance(x, type_conversions.TF_DATASET_REPRESENTATION_TYPES)
self.assertEqual(x.element_spec, tf.TensorSpec(shape=(), dtype=tf.bool))
def test_make_whimsy_element_for_type_spec_raises_SequenceType(self):
type_spec = computation_types.SequenceType(tf.float32)
with self.assertRaisesRegex(ValueError,
'Cannot construct array for TFF type'):
tensorflow_utils.make_whimsy_element_for_type_spec(type_spec)
def serialize_value(value, type_spec=None):
"""Serializes a value into `executor_pb2.Value`.
Args:
value: A value to be serialized.
type_spec: Optional type spec, a `tff.Type` or something convertible to it.
Returns:
A tuple `(value_proto, ret_type_spec)` where `value_proto` is an instance
of `executor_pb2.Value` with the serialized content of `value`, and the
returned `ret_type_spec` is an instance of `tff.Type` that represents the
TFF type of the serialized value.
Raises:
TypeError: If the arguments are of the wrong types.
ValueError: If the value is malformed.
"""
type_spec = computation_types.to_type(type_spec)
if isinstance(value, computation_pb2.Computation):
type_spec = type_utils.reconcile_value_type_with_type_spec(
type_serialization.deserialize_type(value.type), type_spec)
return executor_pb2.Value(computation=value), type_spec
elif isinstance(value, computation_impl.ComputationImpl):
return serialize_value(
computation_impl.ComputationImpl.get_proto(value),
type_utils.reconcile_value_with_type_spec(value, type_spec))
elif isinstance(type_spec, computation_types.TensorType):
return serialize_tensor_value(value, type_spec)
elif isinstance(type_spec, computation_types.NamedTupleType):
type_elements = anonymous_tuple.to_elements(type_spec)
val_elements = anonymous_tuple.to_elements(
anonymous_tuple.from_container(value))
tup_elems = []
for (e_name, e_type), (_, e_val) in zip(type_elements, val_elements):
e_proto, _ = serialize_value(e_val, e_type)
tup_elems.append(
executor_pb2.Value.Tuple.Element(
name=e_name if e_name else None, value=e_proto))
result_proto = (executor_pb2.Value(tuple=executor_pb2.Value.Tuple(
element=tup_elems)))
return result_proto, type_spec
elif isinstance(type_spec, computation_types.SequenceType):
if not isinstance(value,
type_conversions.TF_DATASET_REPRESENTATION_TYPES):
raise TypeError(
'Cannot serialize Python type {!s} as TFF type {!s}.'.format(
py_typecheck.type_string(type(value)),
type_spec if type_spec is not None else 'unknown'))
value_type = computation_types.SequenceType(
computation_types.to_type(value.element_spec))
if not type_analysis.is_assignable_from(type_spec, value_type):
raise TypeError(
'Cannot serialize dataset with elements of type {!s} as TFF type {!s}.'
.format(value_type,
type_spec if type_spec is not None else 'unknown'))
return serialize_sequence_value(value), type_spec
elif isinstance(type_spec, computation_types.FederatedType):
if type_spec.all_equal:
value = [value]
else:
py_typecheck.check_type(value, list)
items = []
for v in value:
it, it_type = serialize_value(v, type_spec.member)
type_analysis.check_assignable_from(type_spec.member, it_type)
items.append(it)
result_proto = executor_pb2.Value(
federated=executor_pb2.Value.Federated(
type=type_serialization.serialize_type(type_spec).federated,
value=items))
return result_proto, type_spec
else:
raise ValueError(
'Unable to serialize value with Python type {} and {} TFF type.'.
format(str(py_typecheck.type_string(type(value))),
str(type_spec) if type_spec is not None else 'unknown'))
def test_to_representation_for_type_with_sequence_type(self):
foo = [1, 2, 3]
self.assertEqual(
reference_executor.to_representation_for_type(
foo, computation_types.SequenceType(tf.int32)), foo)
@computations.tf_computation(tf.string)
def float_dataset_computation(x):
del x # Unused
return tf.data.Dataset.range(5, output_type=tf.float32)
@computations.tf_computation(tf.int32)
def int_identity(x):
return x
@computations.federated_computation(
tf.int32,
computation_types.FederatedType(computation_types.SequenceType(tf.int64),
placements.CLIENTS),
tf.float32,
)
def test_int64_sequence_struct_computation(a, dataset, b):
return a, dataset, b
@computations.federated_computation(
computation_types.FederatedType(computation_types.SequenceType(tf.int64),
placements.CLIENTS))
def test_int64_sequence_computation(dataset):
del dataset
return intrinsics.federated_value(5, placements.SERVER)
def test_serialize_sequence_not_a_dataset(self):
with self.assertRaisesRegex(
TypeError,
r'Cannot serialize Python type int as .* float32\*'):
_ = executor_service_utils.serialize_value(
5, computation_types.SequenceType(tf.float32))
def create_preprocess_fn(
num_epochs: int,
batch_size: int,
shuffle_buffer_size: int = NUM_EXAMPLES_PER_CLIENT,
crop_shape: Tuple[int, int, int] = CIFAR_SHAPE,
distort_image=False,
num_parallel_calls: int = tf.data.experimental.AUTOTUNE
) -> computation_base.Computation:
"""Creates a preprocessing function for CIFAR-100 client datasets.
Args:
num_epochs: An integer representing the number of epochs to repeat the
client datasets.
batch_size: An integer representing the batch size on clients.
shuffle_buffer_size: An integer representing the shuffle buffer size on
clients. If set to a number <= 1, no shuffling occurs.
crop_shape: A tuple (crop_height, crop_width, num_channels) specifying the
desired crop shape for pre-processing. This tuple cannot have elements
exceeding (32, 32, 3), element-wise. The element in the last index should
be set to 3 to maintain the RGB image structure of the elements.
distort_image: A boolean indicating whether to perform preprocessing that
includes image distortion, including random crops and flips.
num_parallel_calls: An integer representing the number of parallel calls
used when performing `tf.data.Dataset.map`.
Returns:
A `tff.Computation` performing the preprocessing described above.
Raises:
TypeError: If `crop_shape` is not an iterable.
ValueError: If `num_epochs` is a non-positive integer, if `crop_shape` is
iterable but not length 3.
"""
if num_epochs < 1:
raise ValueError('num_epochs must be a positive integer.')
if not isinstance(crop_shape, collections.abc.Iterable):
raise TypeError('Argument crop_shape must be an iterable.')
crop_shape = tuple(crop_shape)
if len(crop_shape) != 3:
raise ValueError(
'The crop_shape must have length 3, corresponding to a '
'tensor of shape [height, width, channels].')
feature_dtypes = collections.OrderedDict(
coarse_label=computation_types.TensorType(tf.int64),
image=computation_types.TensorType(tf.uint8, shape=(32, 32, 3)),
label=computation_types.TensorType(tf.int64))
image_map_fn = build_image_map(crop_shape, distort_image)
@computations.tf_computation(computation_types.SequenceType(feature_dtypes)
)
def preprocess_fn(dataset):
if shuffle_buffer_size > 1:
dataset = dataset.shuffle(shuffle_buffer_size)
return (
dataset.repeat(num_epochs)
# We map before batching to ensure that the cropping occurs
# at an image level (eg. we do not perform the same crop on
# every image within a batch)
.map(image_map_fn,
num_parallel_calls=num_parallel_calls).batch(batch_size))
return preprocess_fn
def infer_type(arg: Any) -> Optional[computation_types.Type]:
"""Infers the TFF type of the argument (a `computation_types.Type` instance).
WARNING: This function is only partially implemented.
The kinds of arguments that are currently correctly recognized:
- tensors, variables, and data sets,
- things that are convertible to tensors (including numpy arrays, builtin
types, as well as lists and tuples of any of the above, etc.),
- nested lists, tuples, namedtuples, anonymous tuples, dict, and OrderedDicts.
Args:
arg: The argument, the TFF type of which to infer.
Returns:
Either an instance of `computation_types.Type`, or `None` if the argument is
`None`.
"""
# TODO(b/113112885): Implement the remaining cases here on the need basis.
if arg is None:
return None
elif isinstance(arg, typed_object.TypedObject):
return arg.type_signature
elif tf.is_tensor(arg):
return computation_types.TensorType(arg.dtype.base_dtype, arg.shape)
elif isinstance(arg, TF_DATASET_REPRESENTATION_TYPES):
element_type = computation_types.to_type(arg.element_spec)
return computation_types.SequenceType(element_type)
elif isinstance(arg, structure.Struct):
return computation_types.StructType([
(k, infer_type(v)) if k else infer_type(v)
for k, v in structure.iter_elements(arg)
])
elif py_typecheck.is_attrs(arg):
items = attr.asdict(arg,
dict_factory=collections.OrderedDict,
recurse=False)
return computation_types.StructWithPythonType(
[(k, infer_type(v)) for k, v in items.items()], type(arg))
elif py_typecheck.is_named_tuple(arg):
# In Python 3.8 and later `_asdict` no longer return OrdereDict, rather a
# regular `dict`.
items = collections.OrderedDict(arg._asdict())
return computation_types.StructWithPythonType(
[(k, infer_type(v)) for k, v in items.items()], type(arg))
elif isinstance(arg, dict):
if isinstance(arg, collections.OrderedDict):
items = arg.items()
else:
items = sorted(arg.items())
return computation_types.StructWithPythonType([(k, infer_type(v))
for k, v in items],
type(arg))
elif isinstance(arg, (tuple, list)):
elements = []
all_elements_named = True
for element in arg:
all_elements_named &= py_typecheck.is_name_value_pair(element)
elements.append(infer_type(element))
# If this is a tuple of (name, value) pairs, the caller most likely intended
# this to be a StructType, so we avoid storing the Python container.
if elements and all_elements_named:
return computation_types.StructType(elements)
else:
return computation_types.StructWithPythonType(elements, type(arg))
elif isinstance(arg, str):
return computation_types.TensorType(tf.string)
elif isinstance(arg, (np.generic, np.ndarray)):
return computation_types.TensorType(tf.dtypes.as_dtype(arg.dtype),
arg.shape)
else:
dtype = {
bool: tf.bool,
int: tf.int32,
float: tf.float32
}.get(type(arg))
if dtype:
return computation_types.TensorType(dtype)
else:
# Now fall back onto the heavier-weight processing, as all else failed.
# Use make_tensor_proto() to make sure to handle it consistently with
# how TensorFlow is handling values (e.g., recognizing int as int32, as
# opposed to int64 as in NumPy).
try:
# TODO(b/113112885): Find something more lightweight we could use here.
tensor_proto = tf.make_tensor_proto(arg)
return computation_types.TensorType(
tf.dtypes.as_dtype(tensor_proto.dtype),
tf.TensorShape(tensor_proto.tensor_shape))
except TypeError as err:
raise TypeError(
'Could not infer the TFF type of {}: {}'.format(
py_typecheck.type_string(type(arg)), err))
computation_types.FunctionType(parameter=type_constructors.at_server(
computation_types.AbstractType('T')),
result=type_constructors.at_clients(
computation_types.AbstractType('T'),
True)))
# Materializes client items as a sequence on the server.
#
# Type signature: {T}@CLIENTS -> T*@SERVER
FEDERATED_COLLECT = IntrinsicDef(
'FEDERATED_COLLECT', 'federated_collect',
computation_types.FunctionType(
parameter=type_constructors.at_clients(
computation_types.AbstractType('T')),
result=type_constructors.at_server(
computation_types.SequenceType(
computation_types.AbstractType('T')))))
# Maps member constituents of a client value pointwise using a given mapping
# function that operates independently on each client.
#
# Type signature: <(T->U),{T}@CLIENTS> -> {U}@CLIENTS
FEDERATED_MAP = IntrinsicDef(
'FEDERATED_MAP', 'federated_map',
computation_types.FunctionType(parameter=[
computation_types.FunctionType(computation_types.AbstractType('T'),
computation_types.AbstractType('U')),
type_constructors.at_clients(computation_types.AbstractType('T')),
],
result=type_constructors.at_clients(
computation_types.AbstractType('U'))))
