def test_tff_model_from_functional_fails_with_repeated_metric_names(self):
dataset = create_test_dataset()
input_spec = dataset.element_spec
functional_model = functional.FunctionalModel(initial_weights(),
forward_pass,
predict_on_batch, input_spec)
metric_constructors = [
lambda: tf.keras.metrics.MeanSquaredError(name='same_name'),
lambda: tf.keras.metrics.RootMeanSquaredError(name='same_name')
]
with self.assertRaisesRegex(ValueError,
'each metric should have a unique name'):
functional.model_from_functional(functional_model, metric_constructors)
def test_tff_model_from_functional_federated_aggregate_metrics_succeeds(self):
dataset = create_test_dataset()
input_spec = dataset.element_spec
functional_model = functional.FunctionalModel(initial_weights(),
forward_pass,
predict_on_batch, input_spec)
metric_constructors = [
lambda: tf.keras.metrics.MeanSquaredError(name='mse'),
lambda: tf.keras.metrics.MeanAbsoluteError(name='mae')
]
tff_model = functional.model_from_functional(functional_model,
metric_constructors)
client_1_local_outputs = collections.OrderedDict(
loss=[1.0, 2.0], mse=[1.0, 2.0], mae=[1.0, 2.0])
client_2_local_outputs = collections.OrderedDict(
loss=[2.0, 4.0], mse=[2.0, 2.0], mae=[1.0, 6.0])
metrics_aggregator = aggregator.sum_then_finalize
unfinalized_metrics_type = type_conversions.type_from_tensors(
tff_model.report_local_unfinalized_metrics())
metrics_aggregation_computation = metrics_aggregator(
tff_model.metric_finalizers(), unfinalized_metrics_type)
aggregated_metrics = metrics_aggregation_computation(
[client_1_local_outputs, client_2_local_outputs])
self.assertAllClose(
aggregated_metrics,
# loss = (1.0+2.0)/(2.0+4.0) = 0.5
# mse = (1.0+2.0)/(2.0+2.0) = 0.75
# mae = (1.0+1.0)/(2.0+6.0) = 0.25
collections.OrderedDict(loss=0.5, mse=0.75, mae=0.25))
def test_tff_model_from_functional_binding_metrics_succeeds(self):
dataset = create_test_dataset()
input_spec = dataset.element_spec
functional_model = functional.FunctionalModel(initial_weights(),
forward_pass,
predict_on_batch, input_spec)
metric_constructors = [
tf.keras.metrics.MeanSquaredError, tf.keras.metrics.RootMeanSquaredError
]
tff_model = functional.model_from_functional(functional_model,
metric_constructors)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.05)
loss = None
num_epochs = 50
for batch in dataset.repeat(num_epochs):
with tf.GradientTape() as tape:
batch_output = tff_model.forward_pass(batch, training=True)
gradients = tape.gradient(batch_output.loss,
tff_model.trainable_variables)
optimizer.apply_gradients(zip(gradients, tff_model.trainable_variables))
loss = batch_output.loss
# Expect some amount of convergence after a few epochs of the dataset.
self.assertLess(loss, 0.1)
self.assertAllClose(
tff_model.trainable_variables, ([[1.0, 2.0, 3.0]], [5.0]), atol=0.5)
self.assertAllClose(
tff_model.report_local_unfinalized_metrics(),
collections.OrderedDict(
# The model uses mean squred error as `loss`, so the other two
# metrics (`mean_squared_error` and `root_mean_squared_error`)
# should have the same state as `loss`.
loss=[1066.19628, 1250.0],
mean_squared_error=[1066.19628, 1250.0],
root_mean_squared_error=[1066.19628, 1250.0]))
def test_tff_model_from_functional_same_result(self):
dataset = create_test_dataset()
input_spec = dataset.element_spec
functional_model = functional.FunctionalModel(initial_weights(),
forward_pass,
predict_on_batch, input_spec)
tff_model = functional.model_from_functional(functional_model)
for training in [True, False]:
for batch in dataset:
self.assertAllClose(
tff_model.predict_on_batch(batch[0], training),
functional_model.predict_on_batch(functional_model.initial_weights,
batch[0], training))
tf.nest.map_structure(
self.assertAllClose, tff_model.forward_pass(batch, training),
functional_model.forward_pass(functional_model.initial_weights,
batch, training))
def test_tff_model_from_functional_resets_metrics(self):
dataset = create_test_dataset()
input_spec = dataset.element_spec
functional_model = functional.FunctionalModel(initial_weights(),
forward_pass,
predict_on_batch, input_spec)
metric_constructors = [
tf.keras.metrics.MeanSquaredError, tf.keras.metrics.RootMeanSquaredError
]
tff_model = functional.model_from_functional(functional_model,
metric_constructors)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.05)
expected_initial_local_variables = [0.0, 0, 0.0, 0.0, 0.0, 0.0]
self.assertSequenceEqual(tff_model.local_variables,
expected_initial_local_variables)
# Execute forward pass on the dataset, assert metrics are not zero.
for batch in dataset:
with tf.GradientTape() as tape:
batch_output = tff_model.forward_pass(batch, training=True)
gradients = tape.gradient(batch_output.loss,
tff_model.trainable_variables)
optimizer.apply_gradients(zip(gradients, tff_model.trainable_variables))
loss = batch_output.loss
self.assertGreater(loss, 0)
self.assertAllClose(
tff_model.report_local_unfinalized_metrics(),
collections.OrderedDict(
loss=[876.11523, 25.0],
mean_squared_error=[876.11523, 25.0],
root_mean_squared_error=[876.11523, 25.0]))
# Reset metrics variables.
tff_model.reset_metrics()
self.assertSequenceEqual(tff_model.local_variables,
expected_initial_local_variables)
self.assertEqual(
tff_model.report_local_unfinalized_metrics(),
collections.OrderedDict(
loss=[0, 0],
mean_squared_error=[0, 0],
root_mean_squared_error=[0, 0]))
def test_save_load_convert_to_tff_model(self, model_fn):
dataset = get_dataset()
functional_model = model_fn(input_spec=dataset.element_spec)
path = self.get_temp_dir()
serialization.save_functional_model(functional_model, path)
loaded_model = serialization.load_functional_model(path)
tff_model = functional.model_from_functional(loaded_model)
example_batch = next(iter(dataset))
for training in [True, False]:
self.assertAllClose(
tff_model.predict_on_batch(x=example_batch[0],
training=training), [[0.]] * 5)
for training in [True, False]:
tf.nest.map_structure(
lambda x, y: self.assertAllClose(x, y, atol=1e-2, rtol=1e-2),
tff_model.forward_pass(batch_input=example_batch,
training=training),
model_lib.BatchOutput(loss=74.250,
predictions=np.zeros(shape=[5, 1]),
num_examples=5))
def test_tff_model_from_functional_converges(self):
dataset = create_test_dataset()
input_spec = dataset.element_spec
functional_model = functional.FunctionalModel(initial_weights(),
forward_pass,
predict_on_batch, input_spec)
tff_model = functional.model_from_functional(functional_model)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.05)
loss = None
num_epochs = 50
for batch in dataset.repeat(num_epochs):
with tf.GradientTape() as tape:
batch_output = tff_model.forward_pass(batch, training=True)
gradients = tape.gradient(batch_output.loss,
tff_model.trainable_variables)
optimizer.apply_gradients(zip(gradients, tff_model.trainable_variables))
loss = batch_output.loss
# Expect some amount of convergence after a few epochs of the dataset.
self.assertLess(loss, 0.1)
self.assertAllClose(
tff_model.trainable_variables, ([[1.0, 2.0, 3.0]], [5.0]), atol=0.5)
self.assertAllClose(tff_model.report_local_unfinalized_metrics(),
collections.OrderedDict(loss=[1066.19628, 1250.0]))
def _predict_on_batch(dataset):
tff_model = functional.model_from_functional(loaded_model)
example_batch = get_example_batch(dataset)
return tff_model.predict_on_batch(example_batch[0])