def test_clip_by_global_norm(self): clip_norm = 20.0 aggregate_fn = aggregate_fns.build_clip_norm_aggregate_fn(clip_norm) # Global l2 norms [17.74824, 53.99074]. deltas = [create_weights_delta(), create_weights_delta(constant=10)] deltas_type = tff.framework.type_from_tensors(deltas[0]) weights = [1., 1.] @tff.federated_computation(tff.FederatedType(deltas_type, tff.CLIENTS), tff.FederatedType(tf.float32, tff.CLIENTS)) def federated_aggregate_test(deltas, weights): state = tff.federated_value(aggregate_fn.initialize(), tff.SERVER) return aggregate_fn(state, deltas, weights) federated_aggregate_test.type_signature.result.check_equivalent_to( tff.NamedTupleType(( tff.FederatedType( aggregate_fns.ClipNormAggregateState(clip_norm=tf.float32, max_norm=tf.float32), tff.SERVER), tff.FederatedType(deltas_type, tff.SERVER), ))) state, mean = federated_aggregate_test(deltas, weights) expected_clipped = [] for delta in deltas: flat = tf.nest.flatten(delta) clipped, _ = tf.clip_by_global_norm(flat, clip_norm) expected_clipped.append(tf.nest.pack_sequence_as(delta, clipped)) expected_mean = tf.nest.map_structure(lambda a, b: (a + b) / 2, *expected_clipped) self.assertEqual(state.clip_norm, tf.constant(20.0, tf.float32)) self.assertEqual(state.max_norm, tf.constant(53.99074, tf.float32)) tf.nest.map_structure(self.assertAllEqual, expected_mean, mean)
def build_federated_averaging_process( model_fn, client_optimizer_fn, server_optimizer_fn=lambda: flars_optimizer.FLARSOptimizer(learning_rate= 1.0)): """Builds the TFF computations for optimization using federated averaging. Args: model_fn: A no-arg function that returns a `tff.learning.TrainableModel`. client_optimizer_fn: A no-arg function that returns a `tf.keras.optimizers.Optimizer` for the local client training. server_optimizer_fn: A no-arg function that returns a `tf.keras.optimizers.Optimizer` for applying updates on the server. Returns: A `tff.utils.IterativeProcess`. """ dummy_model_for_metadata = model_fn() type_signature_grads_norm = tff.NamedTupleType([ weight.dtype for weight in tf.nest.flatten( _get_weights(dummy_model_for_metadata).trainable) ]) server_init_tf = build_server_init_fn(model_fn, server_optimizer_fn) server_state_type = server_init_tf.type_signature.result server_update_fn = build_server_update_fn(model_fn, server_optimizer_fn, server_state_type, server_state_type.model, type_signature_grads_norm) tf_dataset_type = tff.SequenceType(dummy_model_for_metadata.input_spec) client_update_fn = build_client_update_fn(model_fn, client_optimizer_fn, tf_dataset_type, server_state_type.model) federated_server_state_type = tff.FederatedType(server_state_type, tff.SERVER) federated_dataset_type = tff.FederatedType(tf_dataset_type, tff.CLIENTS) run_one_round_tff = build_run_one_round_fn(server_update_fn, client_update_fn, dummy_model_for_metadata, federated_server_state_type, federated_dataset_type) return tff.utils.IterativeProcess(initialize_fn=tff.federated_computation( lambda: tff.federated_value(server_init_tf(), tff.SERVER)), next_fn=run_one_round_tff)
# all_samples = [i for i in range(int(num*(len(source[1])/NUM_AGENT)), int((num+1)*(len(source[1])/NUM_AGENT)))] for i in range(0, len(all_samples), BATCH_SIZE): batch_samples = all_samples[i:i + BATCH_SIZE] output_sequence.append({ 'x': np.array([source[0][i].flatten() / 255.0 for i in batch_samples], dtype=np.float32), 'y': np.array([source[1][i] for i in batch_samples], dtype=np.int32) }) return output_sequence BATCH_TYPE = tff.NamedTupleType([('x', tff.TensorType(tf.float32, [None, 784])), ('y', tff.TensorType(tf.int32, [None]))]) MODEL_TYPE = tff.NamedTupleType([('weights', tff.TensorType(tf.float32, [784, 10])), ('bias', tff.TensorType(tf.float32, [10]))]) @tff.tf_computation(MODEL_TYPE, BATCH_TYPE) def batch_loss(model, batch): predicted_y = tf.nn.softmax(tf.matmul(batch.x, model.weights) + model.bias) return -tf.reduce_mean( tf.reduce_sum(tf.one_hot(batch.y, 10) * tf.log(predicted_y), axis=[1])) @tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def create_train_dataset_preprocess_fn(vocab: List[str], client_batch_size: int, client_epochs_per_round: int, max_seq_len: int, max_training_elements_per_user: int, max_batches_per_user=-1, max_shuffle_buffer_size=10000): """Creates preprocessing functions for stackoverflow data. This function returns a function which takes a dataset and returns a dataset, generally for mapping over a set of unprocessed client datasets during training. Args: vocab: Vocabulary which defines the embedding. client_batch_size: Integer representing batch size to use on the clients. client_epochs_per_round: Number of epochs for which to repeat train client dataset. max_seq_len: Integer determining shape of padded batches. Sequences will be padded up to this length, and sentences longer than `max_seq_len` will be truncated to this length. max_training_elements_per_user: Integer controlling the maximum number of elements to take per user. If -1, takes all elements for each user. max_batches_per_user: If set to a positive integer, the maximum number of batches in each client's dataset. max_shuffle_buffer_size: Maximum shuffle buffer size. Returns: Two functions, the first `preprocess_train` and the second `preprocess_val_and_test`, as described above. """ if client_batch_size <= 0: raise ValueError( 'client_batch_size must be a positive integer; you have ' 'passed {}'.format(client_batch_size)) elif client_epochs_per_round == -1 and max_batches_per_user == -1: raise ValueError( 'Argument client_epochs_per_round is set to -1. If this is' ' intended, then max_batches_per_user must be set to ' 'some positive integer.') elif max_seq_len <= 0: raise ValueError('max_seq_len must be a positive integer; you have ' 'passed {}'.format(max_seq_len)) elif max_training_elements_per_user < -1: raise ValueError( 'max_training_elements_per_user must be an integer at ' 'least -1; you have passed {}'.format( max_training_elements_per_user)) if (max_training_elements_per_user == -1 or max_training_elements_per_user > max_shuffle_buffer_size): shuffle_buffer_size = max_shuffle_buffer_size else: shuffle_buffer_size = max_training_elements_per_user feature_dtypes = [ ('creation_date', tf.string), ('title', tf.string), ('score', tf.int64), ('tags', tf.string), ('tokens', tf.string), ('type', tf.string), ] @tff.tf_computation( tff.SequenceType( tff.NamedTupleType([(name, tff.TensorType(dtype=dtype, shape=())) for name, dtype in feature_dtypes]))) def preprocess_train(dataset): to_ids = build_to_ids_fn(vocab, max_seq_len) dataset = dataset.take(max_training_elements_per_user) if shuffle_buffer_size > 0: logging.info('Adding shuffle with buffer size: %d', shuffle_buffer_size) dataset = dataset.shuffle(shuffle_buffer_size) dataset = dataset.repeat(client_epochs_per_round) dataset = dataset.map(to_ids, num_parallel_calls=tf.data.experimental.AUTOTUNE) dataset = batch_and_split(dataset, max_seq_len, client_batch_size) return dataset.take(max_batches_per_user) return preprocess_train
def get_temperature_sensor_example(): """Constructs `canonical_form.CanonicalForm` for temperature sensors example. The temperature sensor example computes the fraction of sensors that report temperatures over the threshold. Returns: An instance of `canonical_form.CanonicalForm`. """ @tff.tf_computation def initialize(): return {'num_rounds': tf.constant(0)} # The state of the server is a singleton tuple containing just the integer # counter `num_rounds`. server_state_type = tff.NamedTupleType([('num_rounds', tf.int32)]) @tff.tf_computation(server_state_type) def prepare(state): return {'max_temperature': 32.0 + tf.cast(state.num_rounds, tf. float32)} # The initial state of the client is a singleton tuple containing a single # float `max_temperature`, which is the threshold received from the server. client_state_type = tff.NamedTupleType([('max_temperature', tf.float32)]) # The client data is a sequence of floats. client_data_type = tff.SequenceType(tf.float32) @tff.tf_computation(client_data_type, client_state_type) def work(data, state): reduce_result = data.reduce( {'num': np.int32(0), 'max': np.float32(-459.67)}, lambda s, x: {'num': s['num'] + 1, 'max': tf.maximum(s['max'], x)}) return ({'is_over': reduce_result['max'] > state.max_temperature}, {'num_readings': reduce_result['num']}) # The client update is a singleton tuple with a Boolean-typed `is_over`. client_update_type = tff.NamedTupleType([('is_over', tf.bool)]) # The accumulator for client updates is a pair of counters, one for the # number of clients over threshold, and the other for the total number of # client updates processed so far. accumulator_type = tff.NamedTupleType([ ('num_total', tf.int32), ('num_over', tf.int32)]) @tff.tf_computation def zero(): return collections.OrderedDict([ ('num_total', tf.constant(0)), ('num_over', tf.constant(0))]) @tff.tf_computation(accumulator_type, client_update_type) def accumulate(accumulator, update): return collections.OrderedDict([ ('num_total', accumulator.num_total + 1), ('num_over', accumulator.num_over + tf.cast(update.is_over, tf.int32))]) @tff.tf_computation(accumulator_type, accumulator_type) def merge(accumulator1, accumulator2): return collections.OrderedDict([ ('num_total', accumulator1.num_total + accumulator2.num_total), ('num_over', accumulator1.num_over + accumulator2.num_over)]) @tff.tf_computation(merge.type_signature.result) def report(accumulator): return {'ratio_over_threshold': ( tf.cast(accumulator['num_over'], tf.float32) / tf.cast(accumulator['num_total'], tf.float32))} # The type of the combined update is a singleton tuple containing a float # named `ratio_over_threshold`. combined_update_type = tff.NamedTupleType( [('ratio_over_threshold', tf.float32)]) @tff.tf_computation(server_state_type, combined_update_type) def update(state, update): return ({'num_rounds': state.num_rounds + 1}, update) return tff.backends.mapreduce.CanonicalForm(initialize, prepare, work, zero, accumulate, merge, report, update)
return "Hello, World!" print(hello_word()) # %% federated_float_on_clients = tff.FederatedType(tf.float32, tff.CLIENTS) print(str(federated_float_on_clients.member)) print(str(federated_float_on_clients.placement)) print(str(federated_float_on_clients)) print(federated_float_on_clients.all_equal) print(tff.FederatedType(tf.float32, tff.CLIENTS, all_equal=True)) # %% simple_regression_model_type = ( tff.NamedTupleType([('a', tf.float32), ('b', tf.float32)]) ) print(str(simple_regression_model_type)) print(str(tff.FederatedType(simple_regression_model_type, tff.CLIENTS, all_equal=True))) # %% @tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS)) def get_average_temperature(sensor_readings): # print('Getting traced, the argument is "{}".'.format(type(sensor_readings).__name__)) return tff.federated_mean(sensor_readings) print(str(get_average_temperature.type_signature))