[联邦学习TFF]构建自己的联邦学习模型

TFF 全称 tensorflow_federated，为谷歌的联邦学习框架。

在TFF官网的Building Your Own Federated Learning Algorithm
界面中，介绍了如何尽可能多的利用现有的TensorFlow代码，构建一个TFF的模型。

本文的行文结构

本文共分为3个章节，其中第1章介绍了TFF的框架，然后给出了客户端和服务器的模型参数更新函数。第2章到主要介绍Federated core的内容。第3章主要把前2章的内容串起来，构建自己的TFF框架。

1.TFF框架的构成

1.1 TFF框架可以分成4个步骤：

1. 服务器向客户端（server-to-client）传递初始模型参数
2. 客户端更新模型参数
3. 客户端向服务器（client-to-server）传递参数
4. 服务器更新参数

1.2 这4个步骤又可以根据是否使用纯TensorFlow的代码分为两类：

第一类：全部使用TensorFlow代码构建
包括第2和第4步，客户端更新模型参数，服务器更新参数。

第二类：使用Federated Core代码构建
包括第1和第3步，server-to-client和client-to-server

下面就详细介绍纯TensorFlow环节的要点。需要Federated Core构建的放在后面第5部分。

1.2.1 客户端更新参数

可以分成两步：
（1）从服务器模型获取客户端模型参数，此处的服务器模型是经过第一步传递过来的模型
（2）客户端模型在客户端数据集上训练和更新参数

@tf.function
def client_update(model, dataset, server_weights, client_optimizer):
  """Performs training (using the server model weights) on the client's dataset."""
  # Initialize the client model with the current server weights.
  client_weights = model.trainable_variables
  # Assign the server weights to the client model.
  tf.nest.map_structure(lambda x, y: x.assign(y),
                        client_weights, server_weights)

  # Use the client_optimizer to update the local model.
  for batch in dataset:
    with tf.GradientTape() as tape:
      # Compute a forward pass on the batch of data
      outputs = model.forward_pass(batch)

    # Compute the corresponding gradient
    grads = tape.gradient(outputs.loss, client_weights)
    grads_and_vars = zip(grads, client_weights)

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

  return client_weights

1.2.2 服务器更新参数

在服务器上更新参数主要涉及更新策略。此处采用了较为简单的"vanilla" 联合平均算法，直接取各个客户端模型参数的平均值作为服务器的模型参数。这里的模型参数只包括可训练的参数。

@tf.function
def server_update(model, mean_client_weights):