def test_two_param_next_raises(self): @federated_computation.federated_computation(SERVER_INT, MODEL_WEIGHTS_TYPE) def next_fn(state, weights): return MeasuredProcessOutput(state, weights, server_zero()) with self.assertRaises(errors.TemplateNextFnNumArgsError): finalizers.FinalizerProcess(test_initialize_fn, next_fn)
def test_next_return_tuple_raises(self): @federated_computation.federated_computation(SERVER_INT, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def tuple_next_fn(state, weights, update): return state, test_finalizer_result(weights, update), server_zero() with self.assertRaises(errors.TemplateNotMeasuredProcessOutputError): finalizers.FinalizerProcess(test_initialize_fn, tuple_next_fn)
def test_non_server_placed_next_measurements_raises(self): @federated_computation.federated_computation(SERVER_INT, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def next_fn(state, weights, update): return MeasuredProcessOutput( state, test_finalizer_result(weights, update), intrinsics.federated_value(1.0, placements.CLIENTS)) with self.assertRaises(errors.TemplatePlacementError): finalizers.FinalizerProcess(test_initialize_fn, next_fn)
def test_next_return_odict_raises(self): @federated_computation.federated_computation(SERVER_INT, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def odict_next_fn(state, weights, update): return collections.OrderedDict( state=state, result=test_finalizer_result(weights, update), measurements=server_zero()) with self.assertRaises(errors.TemplateNotMeasuredProcessOutputError): finalizers.FinalizerProcess(test_initialize_fn, odict_next_fn)
def test_init_tuple_of_federated_types_raises(self): initialize_fn = federated_computation.federated_computation()( lambda: (server_zero(), server_zero())) @federated_computation.federated_computation( initialize_fn.type_signature.result, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def next_fn(state, weights, update): return MeasuredProcessOutput(state, test_finalizer_result(weights, update), server_zero()) with self.assertRaises(errors.TemplateNotFederatedError): finalizers.FinalizerProcess(initialize_fn, next_fn)
def test_non_federated_init_next_raises(self): initialize_fn = tensorflow_computation.tf_computation(lambda: 0) @tensorflow_computation.tf_computation( tf.int32, computation_types.to_type(model_utils.ModelWeights(tf.float32, ())), tf.float32) def next_fn(state, weights, update): new_weigths = model_utils.ModelWeights(weights.trainable + update, ()) return MeasuredProcessOutput(state, new_weigths, 0) with self.assertRaises(errors.TemplateNotFederatedError): finalizers.FinalizerProcess(initialize_fn, next_fn)
def test_next_return_namedtuple_raises(self): measured_process_output = collections.namedtuple( 'MeasuredProcessOutput', ['state', 'result', 'measurements']) @federated_computation.federated_computation(SERVER_INT, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def namedtuple_next_fn(state, weights, update): return measured_process_output(state, test_finalizer_result(weights, update), server_zero()) with self.assertRaises(errors.TemplateNotMeasuredProcessOutputError): finalizers.FinalizerProcess(test_initialize_fn, namedtuple_next_fn)
def test_next_state_not_assignable(self): @federated_computation.federated_computation(SERVER_INT, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def float_next_fn(state, weights, update): del state return MeasuredProcessOutput( intrinsics.federated_value(0.0, placements.SERVER), test_finalizer_result(weights, update), intrinsics.federated_value(1, placements.SERVER)) with self.assertRaises(errors.TemplateStateNotAssignableError): finalizers.FinalizerProcess(test_initialize_fn, float_next_fn)
def test_non_server_placed_next_weight_param_raises(self): @federated_computation.federated_computation(SERVER_INT, computation_types.at_clients( MODEL_WEIGHTS_TYPE.member), SERVER_FLOAT) def next_fn(state, weights, update): return MeasuredProcessOutput( state, test_finalizer_result(intrinsics.federated_sum(weights), update), server_zero()) with self.assertRaises(errors.TemplatePlacementError): finalizers.FinalizerProcess(test_initialize_fn, next_fn)
def test_finalizer(): @federated_computation.federated_computation( empty_init_fn.type_signature.result, computation_types.at_server(MODEL_WEIGHTS_TYPE), computation_types.at_server(FLOAT_TYPE)) def next_fn(state, weights, updates): new_weights = intrinsics.federated_map( tensorflow_computation.tf_computation(lambda x, y: x + y), (weights.trainable, updates)) new_weights = intrinsics.federated_zip( model_utils.ModelWeights(new_weights, ())) return measured_process.MeasuredProcessOutput(state, new_weights, empty_at_server()) return finalizers.FinalizerProcess(empty_init_fn, next_fn)
def test_construction_with_empty_state_does_not_raise(self): initialize_fn = federated_computation.federated_computation()( lambda: intrinsics.federated_value((), placements.SERVER)) @federated_computation.federated_computation( initialize_fn.type_signature.result, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def next_fn(state, weights, update): return MeasuredProcessOutput( state, test_finalizer_result(weights, update), intrinsics.federated_value(1, placements.SERVER)) try: finalizers.FinalizerProcess(initialize_fn, next_fn) except: # pylint: disable=bare-except self.fail('Could not construct an FinalizerProcess with empty state.')
def test_result_not_assignable_to_weight_raises(self): bad_cast_fn = tensorflow_computation.tf_computation( lambda x: tf.nest.map_structure(lambda y: tf.cast(y, tf.float64), x)) @federated_computation.federated_computation(SERVER_INT, MODEL_WEIGHTS_TYPE, SERVER_FLOAT) def next_fn(state, weights, update): return MeasuredProcessOutput( state, intrinsics.federated_map(bad_cast_fn, test_finalizer_result(weights, update)), server_zero()) with self.assertRaises(finalizers.FinalizerResultTypeError): finalizers.FinalizerProcess(test_initialize_fn, next_fn)
def test_constructs_with_non_model_weights_parameter(self): non_model_weights_type = computation_types.at_server( computation_types.to_type( collections.OrderedDict(trainable=tf.float32, non_trainable=()))) @federated_computation.federated_computation(SERVER_INT, non_model_weights_type, SERVER_FLOAT) def next_fn(state, weights, update): del update return MeasuredProcessOutput(state, weights, server_zero()) try: finalizers.FinalizerProcess(test_initialize_fn, next_fn) except: # pylint: disable=bare-except self.fail('Could not construct a valid FinalizerProcess.')
def _build_kmeans_finalizer(centroids_type: computation_types.Type, num_centroids: int): """Builds a `tff.learning.templates.FinalizerProcess` for k-means.""" @tensorflow_computation.tf_computation def initialize_weights(): return tf.ones((num_centroids, ), dtype=_WEIGHT_DTYPE) @federated_computation.federated_computation def init_fn(): return intrinsics.federated_eval(initialize_weights, placements.SERVER) weights_type = initialize_weights.type_signature.result @tensorflow_computation.tf_computation(centroids_type, weights_type, centroids_type, weights_type) def server_update_tf(current_centroids, current_weights, new_centroid_sums, new_weights): return _update_centroids(current_centroids, current_weights, new_centroid_sums, new_weights) summed_updates_type = computation_types.at_server( computation_types.to_type((centroids_type, weights_type))) @federated_computation.federated_computation( init_fn.type_signature.result, computation_types.at_server(centroids_type), summed_updates_type) def next_fn(state, current_centroids, summed_updates): new_centroid_sums, new_weights = summed_updates updated_centroids, updated_weights = intrinsics.federated_map( server_update_tf, (current_centroids, state, new_centroid_sums, new_weights)) empty_measurements = intrinsics.federated_value((), placements.SERVER) return measured_process.MeasuredProcessOutput(updated_weights, updated_centroids, empty_measurements) return finalizers.FinalizerProcess(init_fn, next_fn)
def test_init_param_not_empty_raises(self): one_arg_initialize_fn = federated_computation.federated_computation( SERVER_INT)(lambda x: x) with self.assertRaises(errors.TemplateInitFnParamNotEmptyError): finalizers.FinalizerProcess(one_arg_initialize_fn, test_next_fn)
def test_next_not_tff_computation_raises(self): with self.assertRaisesRegex(TypeError, r'Expected .*\.Computation, .*'): finalizers.FinalizerProcess( initialize_fn=test_initialize_fn, next_fn=lambda state, w, u: MeasuredProcessOutput(state, w + u, ()))
def test_init_not_tff_computation_raises(self): with self.assertRaisesRegex(TypeError, r'Expected .*\.Computation, .*'): finalizers.FinalizerProcess(initialize_fn=lambda: 0, next_fn=test_next_fn)
def test_construction_does_not_raise(self): try: finalizers.FinalizerProcess(test_initialize_fn, test_next_fn) except: # pylint: disable=bare-except self.fail('Could not construct a valid FinalizerProcess.')
def test_init_state_not_assignable(self): float_initialize_fn = federated_computation.federated_computation()( lambda: intrinsics.federated_value(0.0, placements.SERVER)) with self.assertRaises(errors.TemplateStateNotAssignableError): finalizers.FinalizerProcess(float_initialize_fn, test_next_fn)