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))
