<jupyter_start><jupyter_text>Implement a version of `server_update` that updates the server weights to be the midpoint of model_weights and mean_client_weights.  **Setting up environment**<jupyter_code>!pip install --quiet --upgrade tensorflow-federated-nightly
!pip install --quiet --upgrade nest-asyncio

import nest_asyncio
nest_asyncio.apply()
import collections
import attr
import functools
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

np.random.seed(0)<jupyter_output><empty_output><jupyter_text>**Dataset**<jupyter_code>emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()
NUM_CLIENTS = 10
BATCH_SIZE = 20

def preprocess(dataset):
  def batch_format_fn(element):
    return (tf.reshape(element['pixels'], [-1, 784]), 
            tf.reshape(element['label'], [-1, 1]))
  return dataset.batch(BATCH_SIZE).map(batch_format_fn)                           # return a (features, label) tuple in a batch of BATCH_SIZE
client_ids = np.random.choice(emnist_train.client_ids, size=NUM_CLIENTS,          # Randomly selets "NUM_CLIENTS" from the list of clients in  without replacement
                              replace=False)                                      # Probability of selection of each element can also be passed as an argument

federated_train_data = [preprocess(emnist_train.create_tf_dataset_for_client(x))  # Creates dataset for the selected clients
  for x in client_ids
]<jupyter_output><empty_output><jupyter_text>**Model**  Defining Keras model<jupyter_code>def create_keras_model():
  return tf.keras.models.Sequential([                                             # Signal hidden layer sequential keras model
      tf.keras.layers.Input(shape=(784,)),
      tf.keras.layers.Dense(10, kernel_initializer='zeros'),
      tf.keras.layers.Softmax(),
  ])<jupyter_output><empty_output><jupyter_text>Wrapping Keras model<jupyter_code>def model_fn():
  keras_model = create_keras_model()
  return tff.learning.from_keras_model(
      keras_model,
      input_spec=federated_train_data[0].element_spec,                            # Specification of input
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])<jupyter_output><empty_output><jupyter_text>**Building Federated Learning Algorithm**  Initialization Function<jupyter_code>@tff.tf_computation
def server_init():
  model = model_fn()
  return model.trainable_variables
@tff.federated_computation
def initialize_fn():
  return tff.federated_value(server_init(), tff.SERVER)
whimsy_model = model_fn()
tf_dataset_type = tff.SequenceType(whimsy_model.input_spec)                       # Dataset type
model_weights_type = server_init.type_signature.result                            # Model weight type<jupyter_output><empty_output><jupyter_text>Client Update<jupyter_code>@tf.function
def client_update(model, dataset, server_weights, client_optimizer):              # Performs training (using the server model weights) on the client's dataset.
  
  client_weights = model.trainable_variables                                      # Initialize the client model with the current server weights.
  
  tf.nest.map_structure(lambda x, y: x.assign(y),                                 # Assign the server weights to the client model by assigning each element of server_weights to client_weights
                        client_weights, server_weights)
                                                                                  # Use the client_optimizer to update the local model.
  for batch in dataset:                                                           # For each bacth in input "Dataset"
    with tf.GradientTape() as tape:
      outputs = model.forward_pass(batch)                                         # Compute a forward pass on the batch of data

    grads = tape.gradient(outputs.loss, client_weights)                           # Compute the corresponding gradient of outputs.loss w.r.t. client_weights
    grads_and_vars = zip(grads, client_weights)                                   # Zips "gradients" and "client_weights"

    client_optimizer.apply_gradients(grads_and_vars)                              # Apply the gradient using a client optimizer.

  return client_weights
@tff.tf_computation(tf_dataset_type, model_weights_type)
def client_update_fn(tf_dataset, server_weights):
  model = model_fn()
  client_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
  return client_update(model, tf_dataset, server_weights, client_optimizer)<jupyter_output><empty_output><jupyter_text>Server Update<jupyter_code>@tf.function
def server_update(model, mean_client_weights):                                    # Updates the server model weights as the average of the client model weights.
  model_weights = model.trainable_variables
  
  tf.nest.map_structure(lambda x, y: x.assign((x+y)/2),                           # Assign the midpoint of model_weights and mean_client_weights to the server model.
                        model_weights, mean_client_weights)                       # tf.nest.map_structure() applies func to each entry in structure and returns a new structure.
  return model_weights
@tff.tf_computation(model_weights_type)
def server_update_fn(mean_client_weights):
  model = model_fn()
  return server_update(model, mean_client_weights)<jupyter_output><empty_output><jupyter_text>Next Function<jupyter_code>federated_server_type = tff.FederatedType(model_weights_type, tff.SERVER)
federated_dataset_type = tff.FederatedType(tf_dataset_type, tff.CLIENTS)
@tff.federated_computation(federated_server_type, federated_dataset_type)
def next_fn(server_weights, federated_dataset):
  
  server_weights_at_client = tff.federated_broadcast(server_weights)              # Broadcast the server weights to the clients.
  
  client_weights = tff.federated_map(                                             # Each client computes their updated weights.
      client_update_fn, (federated_dataset, server_weights_at_client))
  
  mean_client_weights = tff.federated_mean(client_weights)                        # The server averages these updates.

  server_weights = tff.federated_map(server_update_fn, mean_client_weights)       # The server updates its model.

  return server_weights
federated_algorithm = tff.templates.IterativeProcess(
    initialize_fn=initialize_fn,
    next_fn=next_fn
)<jupyter_output><empty_output><jupyter_text>**Evaluating Algorithm**<jupyter_code>central_emnist_test = emnist_test.create_tf_dataset_from_all_clients().take(1000) # Taking only 1000 samples
central_emnist_test = preprocess(central_emnist_test)                             # Preprocessing test dataset<jupyter_output><empty_output><jupyter_text>Evaluation on test dataset<jupyter_code>def evaluate(server_state):
  keras_model = create_keras_model()                                              # Creates Keras mode
  keras_model.compile(                                                            # Configures the model for training
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]  
  )
  keras_model.set_weights(server_state)                                           # Sets the weights of model same as the server_state
  keras_model.evaluate(central_emnist_test)                                       # Returns the loss value & metrics values for the model in test mode.
                                                                                  # Computation is done in batch, if batch size not secified, 32 is default value<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>server_state = federated_algorithm.initialize()
evaluate(server_state)
for round in range(15):
  server_state = federated_algorithm.next(server_state, federated_train_data)
evaluate(server_state)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0729456384323641, 'fraction_numerical': 0.004171934260429835, 'mean_word_length': 3.015736040609137, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '394551', 'license': 'Apache-2.0', 'n_tokens_mistral': 2285, 'n_tokens_neox': 2193, 'n_words': 493, 'path': 'Building Your Own Federated Learning Algorithm/Challenge 1.ipynb', 'repo_name': 'DeepaliKushwaha/TensorFlow-Federated', 'size': '7910'}