tff.learning.Model

See Stable See Nightly

View source on GitHub

Represents a model for use in TensorFlow Federated.

Used in the notebooks

Each Model will work on a set of tf.Variables, and each method should be a computation that can be implemented as a tf.function; this implies the class should essentially be stateless from a Python perspective, as each method will generally only be traced once (per set of arguments) to create the corresponding TensorFlow graph functions. Thus, Model instances should behave as expected in both eager and graph (TF 1.0) usage.

In general, tf.Variables may be either:

• Weights, the variables needed to make predictions with the model.
• Local variables, e.g. to accumulate aggregated metrics across calls to forward_pass.

The weights can be broken down into trainable variables (variables that can and should be trained using gradient-based methods), and non-trainable variables (which could include fixed pre-trained layers, or static model data). These variables are provided via the trainable_variables, non_trainable_variables, and local_variables properties, and must be initialized by the user of the Model.

In federated learning, model weights will generally be provided by the server, and updates to trainable model variables will be sent back to the server. Local variables are not transmitted, and are instead initialized locally on the device, and then used to produce aggregated_outputs which are sent to the server.

All tf.Variables should be introduced in __init__; this could move to a build method more inline with Keras (see https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) in the future.

return {
    'num_examples': tff.federated_sum(local_outputs.num_examples),
    'loss': tff.federated_mean(local_outputs.loss)
}

Methods

forward_pass

View source

@abc.abstractmethod
forward_pass(
    batch_input, training=True
)

Runs the forward pass and returns results.

This method should not modify any variables that are part of the model parameters, that is, variables that influence the predictions. Rather, this is done by the training loop.

However, this method may update aggregated metrics computed across calls to forward_pass; the final values of such metrics can be accessed via aggregated_outputs.

Uses in TFF:

• To implement model evaluation.
• To implement federated gradient descent and other non-Federated-Averaging algorithms, where we want the model to run the forward pass and update metrics, but there is no optimizer (we might only compute gradients on the returned loss).
• To implement Federated Averaging.

report_local_outputs

View source

@abc.abstractmethod
report_local_outputs()

Returns tensors representing values aggregated over forward_pass calls.

In federated learning, the values returned by this method will typically be further aggregated across clients and made available on the server.

This method returns results from aggregating across all previous calls to forward_pass, most typically metrics like accuracy and loss. If needed, we may add a clear_aggregated_outputs method, which would likely just run the initializers on the local_variables.

In general, the tensors returned can be an arbitrary function of all the tf.Variables of this model, not just the local_variables; for example, this could return tensors measuring the total L2 norm of the model (which might have been updated by training).

This method may return arbitrarily shaped tensors, not just scalar metrics. For example, it could return the average feature vector or a count of how many times each feature exceed a certain magnitude.