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)