手写数字识别之联邦学习

 本教程以及Federated Learning API，主要提供给希望将自己的TensorFlow模型插入TFF的用户使用，后者主要将其视为黑匣子。 要更深入地了解TFF以及如何实现自己的联合学习算法，请参阅FC Core API的教程-自定义联合算法第1部分和第2部分。

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

• 首先请运行以下命令以确保正确设置的环境。 如果看不到“Hello world”，请参阅安装指南以获取说明。
  pip install tensorflow_federated

from __future__ import absolute_import, division, print_function


import collections
from six.moves import range
import numpy as np
import tensorflow as tf
from tensorflow.python.keras.optimizer_v2 import gradient_descent
from tensorflow_federated import python as tff

nest = tf.contrib.framework.nest

np.random.seed(0)

tf.compat.v1.enable_v2_behavior()

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

让我们从数据开始。 联合学习需要联合数据集，即来自多个用户的数据集合。 联合数据通常是non-i.i.d.，这带来了一系列独特的挑战。

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

下面我们如何加载它。

#@test {"output": "ignore"}
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()

load_data（）返回的数据集是tff.simulation.ClientData客户端数据，一个允许您枚举用户集的接口，用于构造tf.data.Dataset数据集表示特定用户的数据，并查询单个元素的结构。下面是如何使用此界面浏览数据集的内容。请记住，虽然此接口允许您迭代客户端ID，但这只是模拟数据的一个特性。正如您将很快看到的，联合学习框架不使用客户机标识—它们的唯一目的是允许您选择数据子集进行模拟。

len(emnist_train.client_ids)

3383

emnist_train.output_types, emnist_train.output_shapes

(OrderedDict([(u'label', tf.int32), (u'pixels', tf.float32)]), OrderedDict([(u'label', TensorShape([])), (u'pixels', TensorShape([28, 28]))]))

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

example_element = iter(example_dataset).next()

example_element['label'].numpy()

5

#@test {"output": "ignore"}
from matplotlib import pyplot as plt
plt.imshow(example_element['pixels'].numpy(), cmap='gray', aspect='equal')
plt.grid('off')
_ = plt.show()

由于数据已经是tf.data.Dataset，因此可以使用Dataset转换完成预处理。 在这里，我们将28x28的图像展平为784个元素的数组，将各个示例进行混洗，将它们组织成批，然后将特征从像素和标签重命名为x和y，以用于Keras。 我们还对数据集进行重复操作以运行多个epoch。

NUM_EPOCHS = 10
BATCH_SIZE = 20
SHUFFLE_BUFFER = 500


def preprocess(dataset):

  def element_fn(element):
    return collections.OrderedDict([
        ('x', tf.reshape(element['pixels'], [-1])),
        ('y', tf.reshape(element['label'], [1])),
    ])

  return dataset.repeat(NUM_EPOCHS).map(element_fn).shuffle(
      SHUFFLE_BUFFER).batch(BATCH_SIZE)

来验证一下。

#@test {"output": "ignore"}
preprocessed_example_dataset = preprocess(example_dataset)

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

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([[3], [0], [7], [0], [8], [2], [7], [7], [9], [0], [5], [3], [3], [7], [1], [2], [6], [5], [2], [0]], dtype=int32))])

我们几乎已具备构建联邦数据集的所有部分。

在模拟中将联合数据喂给到TFF的方法之一就是使用简单的Python列表，列表中的每个元素保存单个用户的数据，元素可以是列表还可以是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]

现在，我们如何选择客户？

在典型的联邦训练场景中，我们正在处理大量潜在的用户设备，其中只有一小部分可在给定的时间点进行训练。例如，当客户端设备是移动电话时，仅在插入电源，断开移动网络连接或处于空闲状态时才参与训练的训。

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

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

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

#@test {"output": "ignore"}
NUM_CLIENTS = 3

sample_clients = emnist_train.client_ids[0:NUM_CLIENTS]

federated_train_data = make_federated_data(emnist_train, sample_clients)

len(federated_train_data), federated_train_data[0]

(3, )

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

def create_compiled_keras_model():
  model = tf.keras.models.Sequential([
      tf.keras.layers.Dense(
          10, activation=tf.nn.softmax, kernel_initializer='zeros', input_shape=(784,))])
  
  def loss_fn(y_true, y_pred):
    return tf.reduce_mean(tf.keras.losses.sparse_categorical_crossentropy(
        y_true, y_pred))
 
  model.compile(
      loss=loss_fn,
      optimizer=gradient_descent.SGD(learning_rate=0.02),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
  return model

关于编译的一个重要注解。如下所述，当在联邦平均算法中使用时，优化器仅占总体优化算法的一半，因为它仅用于计算每个客户端上的本地模型更新。该算法的其余部分涉及如何在客户端上平均这些更新，以及如何将它们应用到服务器上的全局模型。特别是，这意味着此处使用的优化器和学习率的选择可能需要与您在标准i.d.d上训练模型所使用的选择不同。数据集。我们建议从常规SGD开始，学习速度可能会比平常小。我们在这里使用的学习率尚未经过仔细调整，可以随时尝试。

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

def model_fn():
  keras_model = create_compiled_keras_model()
  return tff.learning.from_compiled_keras_model(keras_model, sample_batch)

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

#@test {"output": "ignore"}
iterative_process = tff.learning.build_federated_averaging_process(model_fn)

刚才发生了什么？ TFF已经构造了联邦计算对，并将它们打包到一个tff.utils.IterativeProcess中，这些计算可以通过属性initialize和next使用。

简而言之，联邦计算是使用TFF的内部语言编写的程序，可以表示各种联邦算法（您可以在自定义算法教程中找到更多有关此的信息）。在这种情况下，生成并打包到iterative_process中的两个计算将实现联邦平均算法。

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

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

#@test {"output": "ignore"}
str(iterative_process.initialize.type_signature)

'( -> ,non_trainable=<>>,optimizer_state=>@SERVER)'

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

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

state = iterative_process.initialize()

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

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

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

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

#@test {"timeout": 600, "output": "ignore"}
state, metrics = iterative_process.next(state, federated_train_data)
print('round  1, metrics={}'.format(metrics))

round 2, metrics=round 1, metrics=

让我们再进行几轮。 如前所述，通常在这一点上，您将为每个回合从新随机选择的用户样本中选择模拟数据的子集，以模拟现实的部署，在该部署中，用户不断地来去去去，但是在本教程中， 为了演示起见，我们将重复使用相同的用户，以便系统快速收敛。

#@test {"skip": true}
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= round 3, metrics= round 4, metrics= round 5, metrics= round 6, metrics= round 7, metrics= round 8, metrics= round 9, metrics= round 10, metrics=

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

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

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

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

第一步是确定我们将要使用的TensorFlow变量。为了使以下代码更清晰易懂，让我们定义一个数据结构来表示整个集合。这将包括变量，例如我们将要训练的权重和偏差，以及将保存我们将在训练过程中更新的各种累积统计信息和计数器的变量，例如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))

有了模型参数的变量和累积统计信息，我们现在可以定义前向通过方法，该方法可以计算损失，发出预测并更新单批输入数据的累积统计信息，如下所示。

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.log(y), reduction_indices=[1]))
  accuracy = tf.reduce_mean(
      tf.cast(tf.equal(predictions, flat_labels), tf.float32))

  num_examples = tf.to_float(tf.size(batch['y']))

  tf.assign_add(variables.num_examples, num_examples)
  tf.assign_add(variables.loss_sum, loss * num_examples)
  tf.assign_add(variables.accuracy_sum, accuracy * num_examples)

  return loss, predictions

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

在这里，我们仅返回平均损失和准确性以及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 {
    '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)
}

输入指标参数对应于上述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))])

# TODO(b/124777499): Remove `autograph=False` when possible.
@tf.contrib.eager.function(autograph=False)
def forward_pass(self, batch, training=True):
  del training
  loss, predictions = mnist_forward_pass(self._variables, batch)
  return tff.learning.BatchOutput(loss=loss, predictions=predictions)

@tf.contrib.eager.function(autograph=False)
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.contrib.eager.functions，因为这有助于确保TensorFlow可以序列化，并且不需要显式的控件依赖项。
以上对于评估和算法（例如联邦SGD）就足够了。但是，对于联邦平均，我们需要指定模型应如何在每个批次上进行本地训练。

class MnistTrainableModel(MnistModel, tff.learning.TrainableModel):

# TODO(b/124777499): Remove `autograph=False` when possible.
@tf.contrib.eager.defun(autograph=False)
def train_on_batch(self, batch):
  output = self.forward_pass(batch)
  optimizer = tf.train.GradientDescentOptimizer(0.02)
  optimizer.minimize(output.loss, var_list=self.trainable_variables)
  return output

• 使用新模型模拟联邦培训

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

iterative_process = tff.learning.build_federated_averaging_process(
  MnistTrainableModel)

state = iterative_process.initialize()

#@test {"timeout": 600, "output": "ignore"}
state, metrics = iterative_process.next(state, federated_train_data)
print('round  1, metrics={}'.format(metrics))

round 1, metrics=

#@test {"skip": true}
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=
round 3, metrics=
round 4, metrics=
round 5, metrics=
round 6, metrics=
round 7, metrics=
round 8, metrics=
round 9, metrics=
round 10, metrics=

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

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

evaluation = tff.learning.build_federated_evaluation(MnistModel)

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

str(evaluation.type_signature)

‘(<,non_trainable=<>>@SERVER,{*}@CLIENTS> -> )’
此时无需担心细节，只需要注意，它采用以下通用形式，类似于tff.utils.IterativeProcess.next，但有两个重要区别。 首先，我们不返回服务器状态，因为评估不会修改模型或状态的任何其他方面-您可以将其视为无状态。 其次，评估只需要模型，不需要服务器状态的任何其他部分，例如优化器变量，这些部分都可能与培训相关联。

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

#@test {"output": "ignore"}
train_metrics = evaluation(state.model, federated_train_data)

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

#@test {"output": "ignore"}
str(train_metrics)

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

federated_test_data = make_federated_data(emnist_test, sample_clients)

len(federated_test_data), federated_test_data[0]

(3,
)

#@test {"output": "ignore"}
test_metrics = evaluation(state.model, federated_test_data)

#@test {"output": "ignore"}
str(test_metrics)

‘’
本教程到此结束。 我们鼓励您使用参数（例如批处理大小，用户数量，时代，学习率等），修改上面的代码以模拟每轮用户随机样本的训练，并探索其他教程。 我们已经开发了。