def _model_fn():
keras_model = _create_test_cnn_model(only_digits=True)
loss = tf.keras.losses.SparseCategoricalCrossentropy()
input_spec = collections.OrderedDict(x=tf.TensorSpec([None, 28, 28, 1],
tf.float32),
y=tf.TensorSpec([None], tf.int32))
return simple_fedavg_tf.KerasModelWrapper(keras_model, input_spec, loss)
def tff_model_fn():
"""Construct a fully initialized model for use in federated averaging."""
sample_batch = get_sample_batch()
keras_model = create_original_fedavg_cnn_model(only_digits=True)
loss = tf.keras.losses.SparseCategoricalCrossentropy()
return simple_fedavg_tf.KerasModelWrapper(keras_model, sample_batch,
loss)
def tff_model_fn():
"""Constructs a fully initialized model for use in federated averaging."""
keras_model = create_original_fedavg_cnn_model(only_digits=True)
loss = tf.keras.losses.SparseCategoricalCrossentropy()
return simple_fedavg_tf.KerasModelWrapper(keras_model,
test_data.element_spec, loss)
def _model_fn():
keras_model = _create_test_cnn_model(only_digits=True)
loss = tf.keras.losses.SparseCategoricalCrossentropy()
return simple_fedavg_tf.KerasModelWrapper(keras_model,
_create_random_batch(), loss)
