As last week, we will follow this tutorial:
This time we will define the model in a way to have more control on every aspect of the training:
The first step is to identify the TensorFlow variables we're going to work with. In order to make the following code more legible, let's define a data structure to represent the entire set. This will include variables such as weights and bias that we will train, as well as variables that will hold various cumulative statistics and counters we will update during training, such as loss_sum, accuracy_sum, and num_examples.
# Define the variables list into a namedtuble
MnistVariables = collections.namedtuple(
'MnistVariables', 'weights bias num_examples loss_sum accuracy_sum')
Here's a method that creates the variables. For the sake of simplicity, we represent all statistics as tf.float32, as that will eliminate the need for type conversions at a later stage. Wrapping variable initializers as lambdas is a requirement imposed by resource variables.
# Create the mnist variables
def create_mnist_variables():
# Return the namedTuple defined above
return MnistVariables(
# Add the weights
weights=tf.Variable(
lambda: tf.zeros(dtype=tf.float32, shape=(784, 10)),
name='weights',
trainable=True),
# Add the bias
bias=tf.Variable(
lambda: tf.zeros(dtype=tf.float32, shape=(10)),
name='bias',
trainable=True),
# Number of example in the dataset
num_examples=tf.Variable(0.0, name='num_examples', trainable=False),
# Sum of the loss
loss_sum=tf.Variable(0.0, name='loss_sum', trainable=False),
# Sum of the accuracy
accuracy_sum=tf.Variable(0.0, name='accuracy_sum', trainable=False))
With the variables for model parameters and cumulative statistics in place, we can now define the forward pass method that computes loss, emits predictions, and updates the cumulative statistics for a single batch of input data, as follows.
# Define the forward pass
def mnist_forward_pass(variables, batch):
# Predict the output
y = tf.nn.softmax(tf.matmul(batch['x'], variables.weights) + variables.bias)
# Take the prediction with the highest prob
predictions = tf.cast(tf.argmax(y, 1), tf.int32)
# Flat the batch of predicitons
flat_labels = tf.reshape(batch['y'], [-1])
# Compute the cross entropy loss by hand of the batch
loss = -tf.reduce_mean(
tf.reduce_sum(tf.one_hot(flat_labels, 10) * tf.math.log(y), axis=[1]))
# Compute the accuracy of the batch
accuracy = tf.reduce_mean(
tf.cast(tf.equal(predictions, flat_labels), tf.float32))
# Count the size on the batch
num_examples = tf.cast(tf.size(batch['y']), tf.float32)
# Add the computed values to our Variables NamedTuple
variables.num_examples.assign_add(num_examples)
variables.loss_sum.assign_add(loss * num_examples)
variables.accuracy_sum.assign_add(accuracy * num_examples)
return loss, predictions
Next, we define a function that returns a set of local metrics, again using TensorFlow. These are the values (in addition to model updates, which are handled automatically) that are eligible to be aggregated to the server in a federated learning or evaluation process.
Here, we simply return the average loss and accuracy, as well as the num_examples, which we'll need to correctly weight the contributions from different users when computing federated aggregates.
# Get the metrics from the clients
def get_local_mnist_metrics(variables):
return collections.OrderedDict(
num_examples=variables.num_examples,
loss=variables.loss_sum / variables.num_examples,
accuracy=variables.accuracy_sum / variables.num_examples)
Finally, we need to determine how to aggregate the local metrics emitted by each device via get_local_mnist_metrics. This is the only part of the code that isn't written in TensorFlow - it's a federated computation expressed in TFF. Here's what it looks like:
@tff.federated_computation
def aggregate_mnist_metrics_across_clients(metrics):
return collections.OrderedDict(
num_examples=tff.federated_sum(metrics.num_examples),
loss=tff.federated_mean(metrics.loss, metrics.num_examples),
accuracy=tff.federated_mean(metrics.accuracy, metrics.num_examples))
The input metrics argument corresponds to the OrderedDict returned by get_local_mnist_metrics above, but critically the values are no longer tf.Tensors - they are "boxed" as tff.Values, to make it clear you can no longer manipulate them using TensorFlow, but only using TFF's federated operators like tff.federated_mean and tff.federated_sum.