搬运/翻译自官方教程：www.tensorflow.org/federated/t…

注意：本博客的实验完成于2021年1月，若附加的代码无法使用，原因是tf的包有所更新，函数可能也相应的做出了调整。因此无法运行性的话需要参考官方教程文档↑。

一、教程

在本教程中，我们使用经典的MNIST训练集来介绍TFF的联邦学习（FL）的API层：tff.learning。它包含了一组高级接口，可用于执行常见的联邦学习任务，例如TensorFlow中实现的针对用户提供的模型来进行联邦训练。

本教程以及联邦学习API，主要供给希望将自己的TensorFlow模型引入TFF的用户使用，它会将TensorFlow模型视为黑盒。想要更深入地了解TFF以及如何实现自己的联邦学习算法，请参阅FC API的教程——自定义联邦算法的第1部分和第2部分。

有关tff.learning的更多信息，请继续参考文本生成联邦学习教程，该教程除涵盖循环模型外，还演示了如何加载预训练的序列化Keras模型，来细化联邦学习和用Keras进行评估。

前序

在开始之前，请运行以下命令以确保正确设置您的环境。如果看不到问候语，请参阅安装指南以来进行安装。

# tensorflow_federated_nightly还引入了tf_nightly，这可能会导致tensorboard安装重复，从而导致错误。
!pip uninstall --yes tensorboard tb-nightly

!pip install --quiet --upgrade tensorflow_federated_nightly
!pip install --quiet --upgrade nest_asyncio
!pip install --quiet tb-nightly  # or tensorboard, but not both

import nest_asyncio
nest_asyncio.apply()

安装成功测试

import collections

import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

np.random.seed(0)

print(tff.federated_computation(lambda: 'Hello, World!')())

输出结果

b'Hello, World!'

输入数据准备

让我们从数据开始。联邦学习需要联邦数据集，即来自多个用户的数据的集合。联邦数据通常是非独立同分布的（i.i.d.：independently identically distribution），这造成了一系列特别的挑战。

为了便于进行实验，我们在TFF存储库中导入了一些数据集，其中包括MNIST的联邦版本，其中包含已使用Leaf重新处理过的原始NIST数据集的版本，以便数据由数字的原始写者输入。由于每个写者都有独特的风格，因此该数据集表现出联邦数据集预期的non-i.i.d的类型的行为。

MNIST 的原始数据集为 NIST，其中包含 81 万张手写的数字，由 3600多个志愿者提供，目标是建立一个识别数字的 ML 模型。
通过调用 TFF 的 FL API，使用已由 GitHub 上的Leaf项目处理的 NIST 数据集版本来分隔每个数据提供者所写的数字

更多有关MNIST数据集的介绍请移步这里。

这是加载数据集的方式。

emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()

load_data（）返回的emnist_train和emnist_test是tff.simulation.ClientData的实例，它是一个接口，允许您枚举用户集、构造表示特定用户数据的tf.data.Dataset、和查询单个元素的结构。这是使用该接口浏览数据集内容的方法。尽管此接口允许遍历客户端ID，但这只是模拟数据的功能。联邦学习框架不使用客户端身份——它们的唯一目的是允许选择数据的子集进行仿真。mnist数据集离线安装方法请戳这里。

安装成功后，运行查看数据集大小：

len(emnist_train.client_ids)

结果：

查看数据集结构：

emnist_train.element_type_structure

结果：

OrderedDict([('label', TensorSpec(shape=(), dtype=tf.int32, name=None)), ('pixels', TensorSpec(shape=(28, 28), dtype=tf.float32, name=None))])

以字典形式存储，pixels是图片的key，label是对应图片标签的key，后面预处理将他们改为x和y

查看数据集中第一个数据的label：

example_dataset = emnist_train.create_tf_dataset_for_client(
    emnist_train.client_ids[0])

example_element = next(iter(example_dataset))

example_element['label'].numpy()

结果：

查看数据集中id为0的样本的输入值（用像素表示，用plot画出结果）。

from matplotlib import pyplot as plt
plt.imshow(example_element['pixels'].numpy(), cmap='gray', aspect='equal')
plt.grid(False)
plt.show()

结果：

探索联邦数据中的异构性

联邦数据通常是non-i.i.d.的，用户通常根据使用模式会有不同的数据分布。一些客户端可能由于本地缺乏数据而在设备上部署较少的训练样本；而某些客户端会部署充足的训练样本。用EMNIST数据集来探索这种联邦系统中典型的数据异构性概念。要注意的是，只有对我们这种一个仿真环境才可以对客户的数据进行深入分析，所有数据都可以在本地使用。但是在实际的生产联邦环境中，是无法获得任何用户的数据的。

首先，让我们获取一个客户的数据样本，以在一台模拟设备上感受示例。因为我们正在使用的数据集是由唯一写者进行撰写的，所以一个客户的数据表示由一个人的笔迹撰写的0到9的数字样本，模拟了一个用户的唯一“使用模式”。

## Example MNIST digits for one client
figure = plt.figure(figsize=(20, 4))
j = 0

for example in example_dataset.take(40):
  plt.subplot(4, 10, j+1)
  plt.imshow(example['pixels'].numpy(), cmap='gray', aspect='equal')
  plt.axis('off')
  j += 1

结果：

现在，我们直观地看每个客户端上每个MNIST数字标签的样本数量。在联邦环境中，每个客户端上的样本数量可能会有所不同，具体取决于用户的行为。

# Number of examples per layer for a sample of clients
f = plt.figure(figsize=(12, 7))
f.suptitle('Label Counts for a Sample of Clients')
for i in range(6):
  client_dataset = emnist_train.create_tf_dataset_for_client(
      emnist_train.client_ids[i])
  plot_data = collections.defaultdict(list)
  for example in client_dataset:
    # Append counts individually per label to make plots
    # more colorful instead of one color per plot.
    label = example['label'].numpy()
    plot_data[label].append(label)
  plt.subplot(2, 3, i+1)
  plt.title('Client {}'.format(i))
  for j in range(10):
    plt.hist(
        plot_data[j],
        density=False,
        bins=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

结果：

现在，让我们可视化每个MNIST标签的每个客户端的平均图像。下面的代码将为一个标签的所有用户样本生成每个像素值的平均值。我们将看到，由于每个人独特的手写风格，就同一个数字而言，一位客户的平均图像与另一位客户的截然不同。

# Each client has different mean images, meaning each client will be nudging
# the model in their own directions locally.

for i in range(5):
  client_dataset = emnist_train.create_tf_dataset_for_client(
      emnist_train.client_ids[i])
  plot_data = collections.defaultdict(list)
  for example in client_dataset:
    plot_data[example['label'].numpy()].append(example['pixels'].numpy())
  f = plt.figure(i, figsize=(12, 5))
  f.suptitle("Client #{}'s Mean Image Per Label".format(i))
  for j in range(10):
    mean_img = np.mean(plot_data[j], 0)
    plt.subplot(2, 5, j+1)
    plt.imshow(mean_img.reshape((28, 28)))
    plt.axis('off')

结果：

用户数据可能夹杂噪声，并且标签可能不可靠。例如，客户2的数据中，我们可以看到对于标签2，可能存在一些标签错误的样本，从而产生了更有噪声的平均图像。

当我们在本地回合中用该用户自己的唯一数据进行学习时，我们可以思考在每个客户端上每个本地训练回合将如何在不同的方向推动模型。在本教程的后面，我们将看到如何从所有客户端获取对模型的每次更新（local model），并将它们汇总到我们新的全局模型中（global model），这个全局模型就表现为从我们每个客户端自己的独特数据中都学到了东西。

预处理输入数据

由于数据已经是tf.data.Dataset，因此可以使用Dataset转换完成预处理。在这里，我们做以下操作：

将28x28的图像展开为784个元素的数组；
将各个样本进行打乱；
将它们组织成batchs；
将特征从pixels和label重命名为x和y，以用于Keras。
对数据集进行重复操作以运行多个回合。

NUM_CLIENTS = 10
NUM_EPOCHS = 5
BATCH_SIZE = 20
SHUFFLE_BUFFER = 100
PREFETCH_BUFFER = 10

def preprocess(dataset):

  def batch_format_fn(element):
    """Flatten a batch `pixels` and return the features as an `OrderedDict`."""
    return collections.OrderedDict(
        x=tf.reshape(element['pixels'], [-1, 784]),
        y=tf.reshape(element['label'], [-1, 1]))

  return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(
      BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)

sample_batch输出如下：每个batch_sample含20对数据

preprocessed_example_dataset = preprocess(example_dataset)
sample_batch = tf.nest.map_structure(lambda x: x.numpy(), next(iter(preprocessed_example_dataset)))
sample_batch

结果：

OrderedDict([('x', array([[1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.],
       ...,
       [1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.]], dtype=float32)), ('y', array([[0],
       [5],
       [0],
       [1],
       [3],
       [0],
       [5],
       [4],
       [1],
       [7],
       [0],
       [4],
       [0],
       [1],
       [7],
       [2],
       [2],
       [0],
       [7],
       [1]], dtype=int32))])

现在已具备构建联邦数据集的所有模块。

选择用户并生成对应用户的数据集

在仿真中向TFF提供联邦数据的一种方法就是简单地将其作为一个Python列表，无论是作为列表还是作为tf.data.Dataset，列表的每个元素都包含单个用户的数据。因为我们已经有一个提供tf.data.Dataset的接口，所以让我们使用它。

下面这是一个简单的辅助函数，它从给定的用户集中构造数据集列表，作为一轮训练或评估的输入。

def make_federated_data(client_data, client_ids):
  return [
      preprocess(client_data.create_tf_dataset_for_client(x))
      for x in client_ids
  ]

如何选择客户？

在典型的联邦训练场景中，我们将会处理大量潜在的用户设备，其中只有一小部分可以在给定的时间进行训练。例如，当客户端设备是移动电话，它们仅在插入电源、已关闭计量网络且处于空闲状态时才参与训练时，就会出现这种情况。

当然，我们处于仿真环境中，所有数据都在本地可用。然后，通常在运行仿真时，我们将从参与每一轮训练的客户子集进行简单地随机抽样，通常每一轮都不同。

就是说，通过研究有关联邦平均算法的论文所发现的那样，在一个在每个回合中具有随机采样的客户子集的系统中实现收敛可能需要一段时间，并且必须在其中进行数百回合的交互是不切实际的。

相反，我们要做的是对一组客户端进行一次采样（取10个client），并在各回合中重复使用同一组客户端，以加快收敛速度（故意过拟合这几位用户的数据）。我们将其作为练习，供读者修改本教程以模拟随机抽样。这相当容易做到（记住，使模型收敛可能要花费一些时间）。

sample_clients = emnist_train.client_ids[0:NUM_CLIENTS]

federated_train_data = make_federated_data(emnist_train, sample_clients)

print('Number of client datasets: {l}'.format(l=len(federated_train_data)))
print('First dataset: {d}'.format(d=federated_train_data[0]))

结果：

Number of client datasets: 10
First dataset: <DatasetV1Adapter shapes: OrderedDict([(x, (None, 784)), (y, (None, 1))]), types: OrderedDict([(x, tf.float32), (y, tf.int32)])>

用Keras创建模型

如果您正在使用Keras，则您可能已经具有构造Keras模型的代码。这是一个满足我们需求的简单模型的示例。

def create_keras_model():
  return tf.keras.models.Sequential([
      tf.keras.layers.Input(shape=(784,)),
      tf.keras.layers.Dense(10, kernel_initializer='zeros'),
      tf.keras.layers.Softmax(),
  ])

注意：我们尚未编译模型。损耗，指标和优化器将在后面介绍。

为了将任何模型与TFF一起使用，需要将其包装在tff.learning.Model接口的实例中，该接口与Keras相似，公开方法用于标记模型的前向过程、元数据属性等，但还引入了其他元素，例如控制计算联邦指标的过程的方式。现在让我们不必担心。如果您具有上面刚刚定义的Keras模型，则可以通过调用tff.learning.from_keras_model让TFF为您包装它，并传递模型和样本数据batch作为参数，如下所示。

def model_fn():
  # We _must_ create a new model here, and _not_ capture it from an external
  # scope. TFF will call this within different graph contexts.
  keras_model = create_keras_model()
  return tff.learning.from_keras_model(
      keras_model,
      input_spec=preprocessed_example_dataset.element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

在联邦数据上训练模型

现在，我们有了一个包装为tff.learning.Model的模型，可以与TFF一起使用，我们可以通过调用辅助函数tff.learning.build_federated_averaging_process来让TFF构造联邦平均算法，如下所示。

参数必须是构造函数（例如上面的model_fn），而不是已经构造的实例，以便可以在TFF控制的上下文中进行模型的构造（如果您对这样做的原因感到好奇，我们建议您阅读有关自定义算法的后续教程）。

关于下面的联邦平均算法的一个重要说明，有2个优化器：_clientoptimizer和_serveroptimizer。 _clientoptimizer仅用于在每个客户端上计算本地模型更新。 _serveroptimizer将平均更新应用于服务器上的全局模型。特别是，这意味着所选择的优化器和学习率可能需要与您在标准i.i.d数据集上训练模型所选择的不同。我们建议从常规SGD开始，学习速度可能会比平常小。我们使用的学习率尚未经过精心调整，可以随时尝试。

iterative_process = tff.learning.build_federated_averaging_process(
    model_fn,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))

刚刚发生了什么？ TFF构建了一对联邦计算，并将它们打包到一个tff.templates.IterativeProcess中，其中，这些计算可以作为一对属性initialize和next使用。

调用 tff.templates.IterativeProcess ，将会返回一个 IterativeProcess 的实例，包含两个函数：initialize() 和 next()

initialize() 用于初始化，返回的是训练开始时的state
next() 输入当前的state，执行一轮计算，得到新的state

作者：Kylin_824 链接：www.jianshu.com/p/d2052d35a… 来源：简书著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。简而言之，联邦计算是使用TFF内部语言编写的程序，可以表示各种联邦算法（您可以在自定义算法教程中找到更多有关此的信息）。在这种情况下，生成并打包到iterative_process中的两个计算将实现联邦平均。

TFF的目标是以一种可以在真正的联邦学习设置中执行计算的方式定义计算，但是目前仅实现了本地执行仿真运行。要在模拟器中执行计算，只需像Python函数一样调用它即可。这个默认的解释环境不是为高性能而设计的，但足以满足本教程的要求。我们希望提供更高性能的仿真运行以促进将来版本中的大规模研究。

让我们从initialize计算开始。与所有联邦计算一样，您可以将其视为一个函数。该计算不接受任何参数，并返回一个结果——服务器上联邦平均过程状态的表示。尽管我们不想深入了解TFF的细节，但了解这种状态看起来可能是有益的。您可以如下所示对其进行可视化。

str(iterative_process.initialize.type_signature)

结果：

'( -> <model=<trainable=<float32[784,10],float32[10]>,non_trainable=<>>,optimizer_state=<int64>,delta_aggregate_state=<value_sum_process=<>,weight_sum_process=<>>,model_broadcast_state=<>>@SERVER)'

虽然上述类型乍看之下似乎有点晦涩难懂，但您可以认识到服务器状态包括一个model（MNIST的初始模型参数，该模型参数将分发到所有设备）和optimizer_state（服务器维护的其他信息，例如用于超参数调度的轮数等）。

让我们调用initialize计算来构造服务器状态。

state = iterative_process.initialize()

联邦计算对中的第二个，next，代表单次联邦平均，包括将服务器状态（包括模型参数）推送给客户端，在设备上对本地数据进行训练，收集并平均模型更新，并在服务器上生成新的更新模型。

从概念上讲，您可以认为next具有如下所示的功能类型。

SERVER_STATE, FEDERATED_DATA -> SERVER_STATE, TRAINING_METRICS

特别是，不应将next()视为在服务器上运行的函数，而应将其视为整个分布式计算的声明性函数表示——即一些输入由服务器（SERVER_STATE）提供，但每个输入设备贡献自己的本地数据集。

让我们进行一次单轮训练并将结果可视化。我们可以将上面已经生成的联邦数据用于用户样本。

state, metrics = iterative_process.next(state, federated_train_data)
print('round  1, metrics={}'.format(metrics))

结果：

round  1, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.11502057), ('loss', 3.244929)]))])

让我们再进行几轮。如前所述，通常在这一点上，将在每个回合从重新随机选择的用户样本中选择一部分数据进行仿真，以仿真一个现实的部署，在这种部署中，用户会不断地来来去去，但是在此交互式场景中，为了演示起见，我们将重复使用相同的用户，以便系统快速收敛。即实际上每一轮都可能选择不同的clients，这里为了简化，重复使用一开始选中的那10个clients

NUM_ROUNDS = 11
for round_num in range(2, NUM_ROUNDS):
  state, metrics = iterative_process.next(state, federated_train_data)
  print('round {:2d}, metrics={}'.format(round_num, metrics))

结果：

round  2, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.14609054), ('loss', 2.9141645)]))])
round  3, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.15205762), ('loss', 2.9237952)]))])
round  4, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.18600823), ('loss', 2.7629454)]))])
round  5, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.20884773), ('loss', 2.622908)]))])
round  6, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.21872428), ('loss', 2.543587)]))])
round  7, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.2372428), ('loss', 2.4210362)]))])
round  8, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.28209877), ('loss', 2.2297976)]))])
round  9, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.2685185), ('loss', 2.195803)]))])
round 10, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('sparse_categorical_accuracy', 0.33868313), ('loss', 2.0523348)]))])

在每轮联邦训练之后，训练损失不断减少，这表明该模型正在收敛。这些训练指标有一些重要的注意事项，请参阅本教程后面的“评估”部分。

在TensorBoard中显示模型指标

接下来，让我们使用Tensorboard来可视化这些联邦计算中的指标。

让我们首先创建目录和相应的摘要编写器以将指标写入其中。

logdir = "/tmp/logs/scalars/training/"
summary_writer = tf.summary.create_file_writer(logdir)
state = iterative_process.initialize()

使用相同的摘要编写器绘制相关的标量指标。

with summary_writer.as_default():
  for round_num in range(1, NUM_ROUNDS):
    state, metrics = iterative_process.next(state, federated_train_data)
    for name, value in metrics['train'].items():
      tf.summary.scalar(name, value, step=round_num)

使用上面指定的根日志目录启动TensorBoard。加载数据可能需要几秒钟。

!ls {logdir}
%tensorboard --logdir {logdir} --port=0

events.out.tfevents.1604020204.isim77-20020ad609500000b02900f40f27a5f6.prod.google.com.686098.10633.v2
events.out.tfevents.1604020602.isim77-20020ad609500000b02900f40f27a5f6.prod.google.com.794554.10607.v2

Launching TensorBoard...
<IPython.core.display.Javascript at 0x7fc5e8d3c128>

# Uncomment and run this this cell to clean your directory of old output for
# future graphs from this directory. We don't run it by default so that if 
# you do a "Runtime > Run all" you don't lose your results.

# !rm -R /tmp/logs/scalars/*

为了以相同方式查看评估指标，您可以创建一个单独的eval文件夹，例如“ logs/scalars/eval”，以写入TensorBoard。

自定义模型

Keras是TensorFlow的推荐高级模型API，我们鼓励在可能的情况下在TFF中使用Keras模型（通过tff.learning.from_keras_model）。

但是，tff.learning提供了一个较低级的模型接口tff.learning.Model，该接口公开了使用模型进行联邦学习所需的最小功能。直接实现此接口（可能仍在使用诸如tf.keras.layers之类的构造块）可以实现最大程度的自定义，而无需修改联邦学习算法的内部结构。

因此，让我们从头开始重新做一遍。

定义模型变量，正向传递和度量

第一步是确定我们将要使用的TensorFlow变量。为了使以下代码更清晰易懂，让我们定义一个数据结构来表示整个集合。这将包括我们将要训练的变量weights和bias，以及将保存我们将在训练中更新的各种累积统计信息和计数器的变量，例如loss_sum，precision_sum和num_examples。

MnistVariables = collections.namedtuple(
    'MnistVariables', 'weights bias num_examples loss_sum accuracy_sum')

然后下面是创建变量的函数。为了简单起见，我们将所有统计信息都表示为tf.float32，因为这样一来，以后就无需进行类型转换了。将变量初始化程序包装为lambda是资源变量强加的要求。

def create_mnist_variables():
  return MnistVariables(
      weights=tf.Variable(
          lambda: tf.zeros(dtype=tf.float32, shape=(784, 10)),
          name='weights',
          trainable=True),
      bias=tf.Variable(
          lambda: tf.zeros(dtype=tf.float32, shape=(10)),
          name='bias',
          trainable=True),
      num_examples=tf.Variable(0.0, name='num_examples', trainable=False),
      loss_sum=tf.Variable(0.0, name='loss_sum', trainable=False),
      accuracy_sum=tf.Variable(0.0, name='accuracy_sum', trainable=False))

有了模型参数的变量和累积统计信息，现在可以定义forward方法，该方法可以计算损失，发出预测并更新单batch输入数据的累积统计信息（根据输入的变量和批大小，得到loss值和预测结果），如下所示。

def mnist_forward_pass(variables, batch):
  y = tf.nn.softmax(tf.matmul(batch['x'], variables.weights) + variables.bias)
  predictions = tf.cast(tf.argmax(y, 1), tf.int32)

  flat_labels = tf.reshape(batch['y'], [-1])
  loss = -tf.reduce_mean(
      tf.reduce_sum(tf.one_hot(flat_labels, 10) * tf.math.log(y), axis=[1]))
  accuracy = tf.reduce_mean(
      tf.cast(tf.equal(predictions, flat_labels), tf.float32))

  num_examples = tf.cast(tf.size(batch['y']), tf.float32)

  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

接下来，我们再次使用TensorFlow定义一个返回一组本地指标的函数。这些除了自动更新的模型更新之外，在联邦学习或评估过程中聚合到服务器的值。

在这里，我们仅返回平均loss，accuracy以及num_examples，在计算联邦聚合时我们需要它们正确地加权不同用户的贡献。

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)

最后，我们需要确定如何通过get_local_mnist_metrics融合每个设备发送的本地指标。这是代码中唯一没有用TensorFlow编写的部分——这是用TFF表达的联邦计算。如果想更深入地学习，请浏览自定义算法教程，但是在大多数应用程序中并不需要。下面显示的模式的变体就足够了：

@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))

输入metrics对应于上述get_local_mnist_metrics返回的OrderedDict，但关键是值不再是tf.Tensors——它们被“包装”为tff.Values，为了清楚起见，不能再使用TensorFlow来操作它们，而只能使用TFF的联邦运算符，例如tff.federated_mean和tff.federated_sum。返回的全局聚合字典定义了一组在服务器上可用的指标。

构造一个tff.learning.Model的实例

完成上述所有操作后，我们就可以构建一个与TFF一起使用的模型，类似于让TFF提取Keras模型时生成的模型。

class MnistModel(tff.learning.Model):

  def __init__(self):
    self._variables = create_mnist_variables()

  @property
  def trainable_variables(self):
    return [self._variables.weights, self._variables.bias]

  @property
  def non_trainable_variables(self):
    return []

  @property
  def local_variables(self):
    return [
        self._variables.num_examples, self._variables.loss_sum,
        self._variables.accuracy_sum
    ]

  @property
  def input_spec(self):
    return collections.OrderedDict(
        x=tf.TensorSpec([None, 784], tf.float32),
        y=tf.TensorSpec([None, 1], tf.int32))

  @tf.function
  def forward_pass(self, batch, training=True):
    del training
    loss, predictions = mnist_forward_pass(self._variables, batch)
    num_exmaples = tf.shape(batch['x'])[0]
    return tff.learning.BatchOutput(
        loss=loss, predictions=predictions, num_examples=num_exmaples)

  @tf.function
  def report_local_outputs(self):
    return get_local_mnist_metrics(self._variables)

  @property
  def federated_output_computation(self):
    return aggregate_mnist_metrics_across_clients

tff.learning.Model定义的抽象方法和属性与上一节介绍变量并定义损失和统计信息的代码段相对应。

这里有几点值得强调：

模型将使用的所有状态都必须为TensorFlow变量，因为TFF在运行时不使用Python（代码应该编写为可以部署到移动设备上；有关详细信息，请参阅自定义算法教程）。
模型应该描述它接受的数据形式（input_spec），通常，TFF是一个强类型的环境，并且想要确定所有组件的类型签名。声明模型输入的格式是其中必不可少的一部分。
尽管从技术上讲不是必需的，但建议将所有TensorFlow逻辑（前向传递，度量计算等）包装为tf.functions，因为这有助于确保TensorFlow可以序列化，并且不需要显式控件依赖项。

以上对于评估和算法（例如联合SGD）就足够了。但是，对于Federated Averaging，我们需要指定模型如何在每个batch上进行本地训练。在构建Federated Averaging算法时，我们将指定一个本地优化器。

使用新模型模拟联邦训练

完成上述所有操作后，其余过程看起来就像我们已经看到的一样——只需用新模型类的构造函数替换模型构造函数，然后在创建的迭代过程中使用两个联合计算来循环训练回合。

iterative_process = tff.learning.build_federated_averaging_process(
    MnistModel,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02))

state = iterative_process.initialize()

state, metrics = iterative_process.next(state, federated_train_data)
print('round  1, metrics={}'.format(metrics))

结果：

round  1, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 3.1527398), ('accuracy', 0.12469136)]))])

for round_num in range(2, 11):
  state, metrics = iterative_process.next(state, federated_train_data)
  print('round {:2d}, metrics={}'.format(round_num, metrics))

结果：

round  2, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.941014), ('accuracy', 0.14218107)]))])
round  3, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.9052832), ('accuracy', 0.14444445)]))])
round  4, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.7491086), ('accuracy', 0.17962962)]))])
round  5, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.5129666), ('accuracy', 0.19526748)]))])
round  6, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.4175923), ('accuracy', 0.23600823)]))])
round  7, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.4273515), ('accuracy', 0.24176955)]))])
round  8, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.2426176), ('accuracy', 0.2802469)]))])
round  9, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.1567981), ('accuracy', 0.295679)]))])
round 10, metrics=OrderedDict([('broadcast', ()), ('aggregation', OrderedDict([('value_sum_process', ()), ('weight_sum_process', ())])), ('train', OrderedDict([('num_examples', 4860.0), ('loss', 2.1092515), ('accuracy', 0.30843621)]))])

要在TensorBoard中查看这些指标，请参考上面“在TensorBoard中显示模型指标”中列出的步骤。

评估

到目前为止，我们所有的实验仅提供了联邦训练指标——整个回合中所有clients训练的所有数据batch的平均指标。这通常就引入了关于过度拟合的问题，特别是因为为简单起见，我们在每一轮中都使用了相同的客户端集，但是在针对Federated Averaging算法的训练指标中还有过度拟合的概念。这很容易看出我们是否想象每个客户端都只有一批数据，并且我们在该批数据上进行了多次迭代（回合）训练。在这种情况下，本地模型将很快完全适合该批次，因此我们平均的本地精度指标将接近1.0。因此，这些训练指标可以被视为训练在进行中的信号，但仅此而已。

要对联邦数据执行评估，可以使用tff.learning.build_federated_evaluation函数构造另一个为此目的而设计的federated computation，并将模型构造函数作为参数传入。请注意，与使用MnistTrainableModel的Federated Averaging不同，传递MnistModel就足够了。评估不执行梯度下降，因此不需要构造优化器。

为了进行实验和研究，当可以使用集中式测试数据集时，文本生成联邦学习演示了另一个评估选择：从联邦学习中获取训练后的权重，将其应用于标准Keras模型，然后在集中数据集上简单地调用tf.keras.models.Model .evaluate()。

evaluation = tff.learning.build_federated_evaluation(MnistModel)

可以按以下方式检查评估函数的抽象类型签名。

str(evaluation.type_signature)

结果：

'(<server_model_weights=<trainable=<float32[784,10],float32[10]>,non_trainable=<>>@SERVER,federated_dataset={<x=float32[?,784],y=int32[?,1]>*}@CLIENTS> -> <num_examples=float32@SERVER,loss=float32@SERVER,accuracy=float32@SERVER>)'

此时无需担心细节，只需注意它采用以下通用形式即可，类似于tff.templates.IterativeProcess.next，但有两个重要区别。首先，我们不返回服务器状态，因为评估不会修改模型或状态的任何其他方面——可以将其视为无状态。其次，评估只需要模型，不需要服务器状态的任何其他部分，例如优化程序变量，这些部分都可能与训练相关联。

SERVER_MODEL, FEDERATED_DATA -> TRAINING_METRICS

让我们对训练期间达到的最新状态进行评估。为了从服务器状态中提取最新的训练模型，您只需访问.model成员，如下所示。

train_metrics = evaluation(state.model, federated_train_data)

这就是我们得到的。请注意，这些数字看起来比上面的最后一轮培训报告的数字略好。按照惯例，迭代训练过程报告的训练指标通常会在训练回合开始时反映模型的性能，因此评估指标将始终领先一步。

str(train_metrics)

结果：

'<num_examples=4860.0,loss=1.7142657041549683,accuracy=0.38683128356933594>'

现在，让我们编译一个联邦数据的测试样本，并对测试数据重新进行评估。数据将来自真实用户的相同样本，但来自截然不同的保留数据集。

federated_test_data = make_federated_data(emnist_test, sample_clients)
len(federated_test_data), federated_test_data[0]

结果：

(10,
 <DatasetV1Adapter shapes: OrderedDict([(x, (None, 784)), (y, (None, 1))]), types: OrderedDict([(x, tf.float32), (y, tf.int32)])>)

test_metrics = evaluation(state.model, federated_test_data)

str(test_metrics)

结果：

'<num_examples=580.0,loss=1.861915111541748,accuracy=0.3362068831920624>'

本教程到此结束。我们鼓励您使用参数（例如批处理大小，用户数量，时期，学习率等），修改上面的代码以模拟每轮用户随机样本的训练。

[Tensorflow]图片分类的联邦学习