0、前言
参考：谷歌联邦学习框架、谷歌联邦学习Blog、Stack Overflow上的讨论
只使用TF实现联邦学习的另两篇博客：【联邦学习】用Tensorflow实现联邦模型AlexNet on CIFAR-10，【联邦学习】用Tensorflow实现联邦学习传递梯度
计划
目前中文社区在联邦学习这一块大多还是科普性的文章，本文希望能从代码实现方面来介绍联邦学习，也同时作为我个人的一份学习笔记，如有描述错误的地方请多多指正，共同学习！
联邦学习概念
联邦学习是谷歌在2016年提出的概念：在分布式的场景下，训练数据分别保存在每个clients中，希望提出一种训练方法：跨多个参与客户端（clients）训练一个共享的全局模型。其中的重点关注的问题包括：
- 参与的clients数量众多，比如每个人的手机、物联网设备等
- clients不能一直保持在线状态，比如只有当手机接入充电器和wifi时才会参与计算
- clients收集的本地数据敏感，比如输入法的输入内容，包含隐私信息，不希望分享给其他人（包括中心节点）
因此，联邦学习希望在保证数据隐私的同时，让众多clients利用自己的数据协同参与训练一个中心模型。以移动设备作为clients为例，大致的训练流程如图：
图中，移动设备作为clients，云服务器作为中心server。训练的三个步骤A、B、C是不断循环迭代的，具体含义为：
- A步骤：clients从中心节点获取当前模型参数（蓝色圆），在本地设备上利用自己的local data对模型进行训练，得到新的模型参数（绿色方块）
- B步骤：多个clients分别训练了不同的模型参数（各种颜色的各种形状），融合为一个模型，比如简单地求平均值。注意到图中各种形状的数量是不等的，可以从两方面理解：1）local data的分布不均，即non-i.i.d.的数据，那么训练出来的模型也分布不均。2）不是所有的clients都参与了训练
- C步骤：中心服务器server将新一轮的模型分发给clients，在这个过程中，server可以有选择性地只让一部分clients参与训练
1、TensorFlow Federated Framework
吹水时间结束，联邦学习的理念其实就这么简单直接。在实现方面，Tensorflow专门为联邦学习推出了一个学习框架（TensorFlow Federated，后文简称TFF），现有的TensorFlow（简称TF）或Keras模型代码通过一些转换后就可以变为联邦学习模型。甚至可以加载单机版的预训练的模型，以迁移学习的模式应用到分散式数据的机器学习中。
阅读本文可能需要：
- 一定的Python技能：函数和类、数据类型、装饰器
- 了解TensorFlow的一些概念：non-eager模式、计算图等
- 了解机器学习的一些概念：模型、数据集、训练、梯度等
就算不了解也没有关系，我将把TFF当作一个新的编程语言（也确实是）来讲解入门，如果哪里写的不清楚请多多留言指出，谢谢！
框架设计理念
在一头扎入技术细节之前，最好先了解TFF框架的设计理念，对于理解各种class的意义有很大帮助。联邦学习的参与角色有：客户端（clients）和服务端（server）。
数据为主
提到clients和server的时候，马上会想到C&S端两套代码、数据交互等分布式的东西。但是在TFF框架中，谷歌并不想让用户去考虑这些东西，它希望用户能够将重点放在数据处理上，而不是代码分离上。因此，在编写TFF代码时，我们不需要指明是某段代码应该运行在Clients端还是Server端（后文简称C端、S端），但是要显式指出每个数据是储存在C端/S端、是全局唯一的还是有多份拷贝的。需要注意，这里提到的“数据”，可以指代：数据集、变量、常量等所有模型会用到的值（官方文档中所称的value）。
回到上图中的例子：蓝色圆和绿色方块是需要学习的目标模型，因此他们是存放在S端、且全局唯一的。将S端的值交给clients去更新时（图中C到A步骤），称为广播（对应tff.federated_broadcast函数，但现在不要去关心具体函数内容）。广播操作会将一个S端的value“转化”成C端的value，而其中的“转化”操作的具体实现对TFF用户来说是隐藏的，不需要去关心。这是TFF框架与分布式训练理念非常不同的一点，需要加倍理解。
同理，在图中的B步骤，把多个C端的value整合成一个S端的value，称为聚合（Aggregation）操作，TFF也提供了许多种预设函数供用户使用，不需要关心数据到底是咋传输的。
整体训练
在编写模型、训练代码的时候，clients和server应当是看作一个整体（也就是“联邦”的含义），不需要分割开S端和C端的代码，完全可以写入同一个文件里，C端和S端的区分是在代码逻辑层面的。例如：一个函数，它既能被C端调用，也能同时被S端调用。但是，某些TFF函数只能接受存放在C端的输入，某些只能接受S端的。
类似TF的non-eager模式一样编写完模型代码和训练代码后，TFF会自动地将代码分别放置到clients和server设备上（但是目前还只能单机模拟哈哈）。我们编写代码的首要且唯一的任务就是训练一个好的模型，只需要关注模型的架构、C&S交互的数据格式、聚合多clients模型的方式就可以了。至于如何协同多设备、怎么在网络上传输都不是TFF用户需要关注的细节。
新的语言
类似TF框架，TFF实际运行的代码并不是Python，而是通过Python代码来编写运算逻辑，实际上是编译成另一种语言去执行。Why？因为许多设备（例如手机、传感器）是很难有Python的运行环境的，更不可能去安装几百Mb的TensorFlow框架，那么在这些设备上执行Python代码的难度是非常大的。因此，编写TFF代码时会遇到非常多“反Python逻辑”的要求，例如：要求提供函数的输入参数的类型、要求使用它规定的几种数据类型、要求确定value的存放位置（S端/C端），原生Python逻辑不会在训练时执行等等。所有的这些，都是为了让TFF编写的代码能编译到low-level，让模型能运行在真实分布式场景下所作出的妥协。当然，随着TFF的更新，某些要求会有所放松也是可能的。
所以，学习TFF框架时，应该把它当成一门新的语言进行学习：这里有新的数据类型、新的变量声明方式、新的函数定义方式、新的执行方法等。原生Python代码和它进行混用时，要特别注意它是否能在训练时生效，是否有预期的结果。一个很简单的例子：在函数定义中print了一些东西，只会在函数定义的时候输出，而在执行的时候没有输出，就是因为函数已经被“编译”成其他语言了，而print没有被编译进去（类似注释）。如果我们编写的一些函数使用到了不同端的数据，在真正执行的时候，一个函数甚至会被拆分到不同的机器上执行。但是，这些情况停留在逻辑设计层面就可以了，编写代码的时候该咋写函数就咋写，还可以嵌套着用，把拆分执行的事情交给TFF去解决。
不同层次的API
又来类比一下：TensorFlow相对于Keras来说，是low-level的接口。编写模型时，TF需要更详细地编写变量之间的运算，而Keras只需要拼接全连接层、卷积层等部件。对应到TFF框架中，也有两个不同层次的接口：
- Federated Learning (FL) API：该层提供了一组高阶接口，使开发者能够将包含的联合训练和评估实现应用于现有的 TensorFlow 模型。
- Federated Core (FC) API：该系统的核心是一组较低阶接口，可以通过在强类型函数式编程环境中结合使用 TensorFlow 与分布式通信运算符，简洁地表达新的联合算法。这一层也是我们构建联合学习的基础。
在本文中，我们将自底向上进行学习，把FC API当作一门新的语言进行学习、实践。
安装TFF库（conda）
详细的安装流程请参见官方指南，这里我们将使用conda进行安装。假设你已经安装好了conda（Anaconda或Miniconda都行），命令行执行以下命令：新建一个名为tf-fed的环境，并安装python3.7：
conda create -n tf-fed python=3.7 --yes
安装完成后，切换到新建的tf-fed环境（以后每次重启命令行都要切换过来）：
conda activate tf-fed
用pip安装TFF，写作时安装版本的是tensorflow_federated==0.13.1（这一步耗时较长，可以去泡杯茶先）：
pip install --upgrade tensorflow_federated
如果下载太慢，可以切换到pip清华源再次安装。
等待安装结束后，验证一下是否安装成功：
python -c "import tensorflow_federated as tff; print(tff.federated_computation(lambda: 'Hello World')())"
如果成功输出了'Hello World'（以及一堆Warning），就说明TFF框架已经安装好了。顺便提一下，本文写作时tff.__version__=0.13.1，它还是一个开发版，因此本文提到的API可能会随着版本更新发生变化。
2020/3/31 更新：初次写作时是0.8.0版，最近陆续有读者反馈无法安装或执行出错，与TFF开发者交流后得知这个版本已经被淘汰了，pip也无法下载到。因此需要安装最新版，代码也需要调整才能执行。
2、数据类型
把TFF当作一门新的编程语言来学习，我们先来了解一下它提供的数据类型有哪些。TFF中，这些数据类型可以分为两类：端无关的，和端有关的，端无关的类型不需要指明存放位置（S端或C端），另一个则需要。定义类型的时候，可以直接print出来看看它的类型和shape，称为compact notation。
为避免误解，请读者先思考这个问题：“类型和变量的关系是什么？”
执行下面的例子前，请先引入库：
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
tf.compat.v1.enable_v2_behavior()
# tff v0.13.1 新版本需要指定默认executor，否则无法eager执行tff.function
tff.framework.set_default_executor(tff.framework.ReferenceExecutor())
端无关的类型
Tensor types（tff.TensorType）
张量类型，需要指定它的元素数据类型dtype和形状shape。其中，dtype的所有可选列表由tf.dtypes.DType指定。例如：
print(tff.TensorType(tf.string, shape=[3, 4, 5]))
print(tff.TensorType(tf.float32, shape=[6, 7]))
print(tff.TensorType(tf.int8, shape=None))
print(tff.TensorType(tf.bool))
将会输出它们的compact notation，代表着类型和形状，而不是输出值（想想numpy.array会输出什么）：
string[3,4,5]
float32[6,7]
int8
bool
Sequence types（tff.SequenceType）
列表类型，其中的元素类型应当为TFF的tff.Type，或者是能转换成tff.Type的东西。例如：
print(tff.SequenceType(tff.TensorType(tf.int8, shape=None)))
print(tff.SequenceType(tf.float32))
print(tff.SequenceType(tff.TensorType(tf.string, shape=[3, 4, 5])))
将会输出：
int8*
float32*
string[3,4,5]*
这里的*就代表着它是一个列表，*前面的是列表中元素的类型。
Named tuple types（tff.NamedTupleType）
如果你用过Python标准库中的collection.namedtuple，那这个类型就是就是TFF框架下的它。没有用过也没有关系，简单来说，tff.NamedTupleType就是元素可以带有key的tuple。看下例子：
import collections
print(tff.NamedTupleType(list((tf.int8, tf.int16))))
print(tff.NamedTupleType(tuple((('key1', tf.float32), (tf.int8), ('key3', tff.SequenceType(tf.string))))))
print(tff.NamedTupleType(collections.OrderedDict([('y', tf.float32), ('x', tff.TensorType(tf.string, shape=[3, 4, 5]))])))
tff.NamedTupleType接受三种类型的输入：list，tuple和collections.OrderedDict（也就是collection.namedtuple生成的subclass生成的对象，一种有序的字典dict类型）
将会输出：
<int8,int16>
<key1=float32,int8,key3=string*>
<y=float32,x=string[3,4,5]>
其中的尖括号是named tuple types的标志。可以注意到，元素的key是可选的，有key无key的元素位置也是很随意的。一般来说，这个类可以用于定义模型的参数、输入输出等。
Function types（tff.FunctionType）
TFF是函数式编程框架，编写的模型是由一个个函数拼接起来的。因此，编写模型代码时，内置函数以及我们编写的函数，作为一个个部件拼凑起来，由TFF后端编译成跨平台的语言进行运算。这些函数具有的类型，就是tff.FunctionType。在编译的过程中，我们需要指定函数的输入类型，且只能有一个输入值，和一个函数返回值（不像python函数那样可以输入多个形参、返回多个值）。
下面举个简单的例子，先不用理解它，后文会详细讲解如何定义函数（也分成端无关和端有关的函数）。
import numpy as np
@tff.tf_computation(tff.SequenceType(tf.int32))
def add_up_integeres(x):
return x.reduce(np.int32(0), lambda x, y: x + y)
print(type(add_up_integeres))
print(add_up_integeres.type_signature)
print(isinstance(add_up_integeres.type_signature, tff.FunctionType))
得到输出：
<class 'tensorflow_federated.python.core.impl.computation_impl.ComputationImpl'>
(int32* -> int32)
True
可以看出，tff.FunctionType不是经过TFF包装的函数本身，而是函数的签名type_signature。它指明了输入输出的形式，用圆括号和箭头->表示。同时，只有一个输入和一个输出，但是输入变量可以是复合类型（例如这里的列表类型），因此是能满足各种需求的。
端有关的类型
端有关的类型，是联邦学习逻辑层面的重点。它们主要完成两件任务：
- 显式地定义数据值应该存放在C端还是S端（
Placement） - 定义这个数据是否全局一致（
All equal?）
我们可以把端有关的类型，想象成一个快递盒：里面装的数据类型是上文提到的端无关类型，盒子上贴着两个标签Placement和All equal?。那么联邦学习的流程，可以类比成：快递盒子从S端中心节点发出，交给C端加工厂处理（盒子由S端传输到C端，还可以复制成多份）。加工厂把盒子里的东西拿出来操作（与端无关）。操作完成后再打包发回中心节点（盒子由C端传输到S端），中心点把各个工厂发回的加工品聚合成一个，就完成了一轮加工。
Placement type
顾名思义，就是定义存放位置的类型。目前，我们能用到的有两个：tff.SERVER和tff.CLIENTS，也就是定义了S端和C端，把他们俩当常数用就好了。
以后也许可以定义新的Placement，以实现更复杂的场景。具体咋用呢，看下面的联邦类型。
Federated types（tff.FederatedType）
以数据驱动的联邦学习，终于到了定义联邦类型的时候了。联邦类型tff.FederatedType把上面提到的端无关类型包装起来，并增加两个属性：
placement（必填），必须是tff.SERVER和tff.CLIENTS这些Placement typeall_equal（可选，默认为None）。类型为bool，代表着这份数据是否全局统一，还是可以有不同的值。如果没有指定all_equal，它会根据placement的值来选择。默认情况下placement=tff.SERVER时all_equal=True，反之为False。
同样，Federated types也可以输出定义，看下面的例子：
import collections
print(tff.FederatedType(tf.int8, tff.CLIENTS))
print(tff.FederatedType(tff.SequenceType(tf.float32), tff.CLIENTS, all_equal=True))
ntt = tff.NamedTupleType(collections.OrderedDict([('y', tf.float32), ('x', tff.TensorType(tf.string, shape=[3, 4, 5]))]))
print(tff.FederatedType(ntt, tff.SERVER))
print(tff.FederatedType(tff.TensorType(tf.int8, shape=None), tff.SERVER, all_equal=False))
将会输出：
{int8}@CLIENTS
float32*@CLIENTS
<y=float32,x=string[3,4,5]>@SERVER
{int8}@SERVER
Federated types的类型表示格式为T@G或{T}@G，其中T为TFF的数据类型，G为存放的位置，花括号{}表示非全局唯一，而没有花括号就表示全局唯一，即all_equal=True。仔细看一下，第1、3个例子是没有指定all_equal，他们是根据placement的值来确定的。
变量声明
解决了类型定义，就可以来声明变量了。每个变量都有一个类型，没错吧？别搞混了类型和变量的概念哈。下面看一下TFF框架里要用到的变量怎么声明：
# 定义一个类型
OUR_TYPE = tff.TensorType(tf.int8, shape=[10])
# 声明一个变量
var = tff.utils.create_variables('var_name', OUR_TYPE)
# 打印一下
print(OUR_TYPE)
print(var)
和TF很类似，上面的代码声明了一个名为var，在计算图中名为var_name，类型为tff.TensorType(tf.int8, shape=[10])的变量。它可以作为模型参数等进行训练和更新。上面的代码会获得如下输出：
int8[10]
<tf.Variable 'var_name:0' shape=(10,) dtype=int8, numpy=array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int8)>
可见，定义的变量实质上是tf.Variable的实例，而不是tff的。那么在实际应用时，如果你对tf比较熟悉，你可以先定义tf的数据类型（或python标准库的类型），再导出tff的类型，如下例子：
# 声明一个tf Tensor类型
ccc = tf.TensorSpec(shape=(10,), name='const_name', dtype=tf.int8)
# 转成类型
OUR_TYPE2 = tff.to_type(ccc)
# 打印一下
print(ccc)
print(OUR_TYPE2)
上面的代码会输出：
TensorSpec(shape=(10,), dtype=tf.int8, name='const_name')
int8[10]
是不是觉得上面的数据类型定义都白学啦？只需要一个函数tff.to_type，里面放进去一个可转换的数据类型，就可以得到tff所需的数据类型。好用是好用，但还是得严谨地学习的嘛。接下来的例子，可能会更多地看到第二种定义类型的方式。
3、函数定义
TFF是函数式编程框架，根据函数是否使用到了不同端的数据来划分，可以分为两种类型：端无关的函数和端有关的函数。
端无关的函数
使用TensorFlow原生计算方式，只需要定义一下输入的类型即可，例如：
import numpy as np
@tff.tf_computation(tff.SequenceType(tf.int32))
def add_up_integeres(x):
return x.reduce(np.int32(0), lambda x, y: x + y)
上面这段代码，使用TFF的APItff.tf_computation来包装一个TensorFlow计算函数，其中显式声明了形参x的类型为tff.SequenceType(tf.int32)。在函数内部，x可以当作tf.data.Dataset来使用成员函数reduce()，这也是上文介绍过的。经过这个包装，函数add_up_integeres()已经被编译成了一个TFF的强类型函数（返回值类型也自动推算出来了），可供其他TFF函数调用，也可以当成普通python函数调用。例如：
print(add_up_integeres.type_signature)
print(add_up_integeres([1,2,3]))
你会得到下面的输出：
(int32* -> int32)
6
注意到两个地方：
- 返回值类型是自动推算出来的（
int32） - 变量的自动类型转换是可以的（
list(int) -> tff.SequenceType(tf.int32)）
端有关的函数
我们也可以自己定义端有关的函数，声明输入形参时，不仅需要指明输入的数据类型，还需要指明它的存放位置。废话不多说，我们看个例子：
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(sensor_readings):
return tff.federated_mean(sensor_readings)
print(get_average_temperature.type_signature)
这里调用了APItff.federated_computation来包装一个端有关的函数，指定了形参sensor_readings的类型为端有关的类型tff.FederatedType(tf.float32, tff.CLIENTS)。在函数中调用了一个TFF内置函数tff.federated_mean，得到函数的返回值。执行上面那段代码，可以得到输出：
{float32}@CLIENTS -> float32@SERVER)
可以看出，输入的数据是C端的浮点数（且all_equal=False），而输出的是S端的浮点数（且all_equal=True）。这里有个细节：我们一般说到“求平均值”都是对一个列表求一个平均值，但是上面定义的函数并不是输入一个列表（tff.SequenceType），而是一个对“一个”值求一个平均值。Why？请同学们结合联邦学习的场景自行理解。
tff.federated_computation的输入类型也可以不是端有关的tff.FederatedType，就像tff.tf_computation那样用，例如：
@tff.federated_computation((tff.FunctionType(tf.int32, tf.int32), tf.int32))
def foo(f, x):
return f(f(x))
print(foo.type_signature)
输出：
(<(int32 -> int32),int32> -> int32)
我们可以这样理解：
tff.tf_computation里定义的是纯TF逻辑tff.federated_computation里可以是TF逻辑（可以调用tff.tf_computation定义的函数），也可以是联邦训练的逻辑（调用其他端有关函数，包括很多TFF内置函数）
内置函数
TFF框架还提供了许多内置函数，例如上面用到的tff.federated_mean。它们也可以分为端有关和端无关的，区分的方法非常简单，如果函数是tff.federated_开头的，就是端有关的，会规定输入输出的存放位置。其他的函数即是端无关的，输入输出的类型也是端无关类型（有些也可以用在端有关场景）。具体有哪些函数可以用，请查阅官方文档：https://www.tensorflow.org/federated/api_docs/python/tff#functions。
其他注意
- 函数调用的语法就和调用普通的python函数一样，但是要注意输入参数的类型是否匹配（或可以自动转换）
tff.tf_computation包装的函数不能调用tff.federated_computation包装的，但反过来可以tff.federated_computation包装的函数输入参数不一定都是Federated Type，但其中的表达的逻辑是端有关的
4、训练模型
学习完数据类型、函数定义之后，我们看一下利用TFF的low-level API如何实现一个简单的逻辑回归模型。这一章节内容主要参考官方文档https://www.tensorflow.org/federated/tutorials/custom_federated_algorithms_2，如果英语阅读能力OK可以去仔细看一看官方文档的细节，本文的讲解会简约很多。
首先我们来引入一些库：
import collections
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
tf.compat.v1.enable_v2_behavior()
# TODO(b/148678573,b/148685415): must use the ReferenceExecutor because it
# supports unbounded references and tff.sequence_* intrinsics.
tff.framework.set_default_executor(tff.framework.ReferenceExecutor())
加载数据集
然后我们加载一下数据集并进行预处理，这里用的是简单的mnist手写数字数据集（第一次加载数据可能需要某些神奇的联网技巧，或提前下载好数据集文件来用，搜一下mnist.npy都会有很多教程）：
mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
NUM_EXAMPLES_PER_USER = 1000
BATCH_SIZE = 100
def get_data_for_digit(source, digit):
output_sequence = []
all_samples = [i for i, d in enumerate(source[1]) if d == digit]
for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):
batch_samples = all_samples[i:i + BATCH_SIZE]
output_sequence.append({
'x':
np.array([source[0][i].flatten() / 255.0 for i in batch_samples],
dtype=np.float32),
'y':
np.array([source[1][i] for i in batch_samples], dtype=np.int32)
})
return output_sequence
federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)]
federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]
mnist数据集是手写数字数据集，每个样本是一张灰阶图片（28x28），代表着0~9之间的某个数字，也就是该样本的标签，所以它是一个10分类数据集。可视化如下图：
上面预处理的时候，做了几件事情：
- 把每个样本从矩阵变成一维向量（长度为
28x28=784），元素从灰度（0~255）缩放到0~1的范围 - 把同一个数字的样本手动挑出来放到一个list里，并且分了batch
- 得到的
federated_train_data是一个列表，长度为10，每个元素分别存放了对应下标的数字样本的batch列表
啰嗦的有点多，其实大家把federated_train_data的长度、元素类型等输出出来看看马上就理解了。
定义batch、model的联邦类型
接下来我们定义TFF框架的类型，准备用到模型算法里。首先定义数据集的输入输出类型，使用的是先定义TF类型，再转换成TFF类型的方法：
BATCH_SPEC = collections.OrderedDict(
x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),
y=tf.TensorSpec(shape=[None], dtype=tf.int32))
BATCH_TYPE = tff.to_type(BATCH_SPEC)
print(str(BATCH_TYPE))
输出：
<x=float32[?,784],y=int32[?]>
BATCH_TYPE是一个tff.NamedTupleType类型，其中有两个key，分别是：
x，类型为tff.TensorType(tf.float32, [None, 784])：一个0维长度不确定、1维长度为784的二维浮点数张量。代表输入的样本，0维长度不确定是因为我们可以一批输入任意多个样本，通过矩阵运算来加速模型训练y，类型为tff.TensorType(tf.int32, [None])：一个0维长度不确定的一维整数张量。代表输出的标签，y的0维长度应当和x的相同，表示样本和标签一一对应。
又啰嗦了一下，毕竟是入门教程嘛。接下来定义我们要学习的模型，这里选用简单的逻辑回归模型，包括两个参数：权重（weights）和偏置（bias）：
MODEL_SPEC = collections.OrderedDict(
weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),
bias=tf.TensorSpec(shape=[10], dtype=tf.float32))
MODEL_TYPE = tff.to_type(MODEL_SPEC)
print(MODEL_TYPE)
输出：
<weights=float32[784,10],bias=float32[10]>
和上面非常类似的定义方式，这里就不啰嗦了。需要注意的是：目前定义的BATCH_TYPE和MODEL_TYPE都是端无关的，还没有指定存放位置。
定义loss函数
有了模型参数类型，我们来定义一下损失函数batch_loss。其中，我们把前向传播的计算过程（forward_pass）提取了出来，单独写作一个函数，便于后续复用。
而且我们注意到，前向传播是用@tf.function来包装的，因为后续我们会在另一个@tf.function里调用forward_pass，而@tf.function不能调用@tff.tf_computation包装的函数，反过来则可以调用（就是下面这个例子）。
@tf.function
def forward_pass(model, batch):
predicted_y = tf.nn.softmax(
tf.matmul(batch['x'], model['weights']) + model['bias'])
return -tf.reduce_mean(
tf.reduce_sum(
tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE)
def batch_loss(model, batch):
return forward_pass(model, batch)
装饰器tff.tf_computation指定了输入参数的类型，而函数的内容就是常见的TF loss定义了。在输入数据batch和模型参数的使用上，其实和Python的dict没什么区别（严格来说是collections.OrderedDict）。所以说，上面定义了一堆TFF的类型，学习模型的主体内容并不会变复杂多少。
我们可以初始化一个值全为零的模型，喂给它一个batch的数据，计算一下loss：
initial_model = collections.OrderedDict(
weights=np.zeros([784, 10], dtype=np.float32),
bias=np.zeros([10], dtype=np.float32))
sample_batch = federated_train_data[5][-1]
print(batch_loss(initial_model, sample_batch))
输出：
2.3025854
定义model优化方法
有了loss函数之后，就可以对模型进行优化对吧？那么接下来，我们定义一个函数，输入一个batch，用TF的随机梯度下降优化器来优化上面那个loss：
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def batch_train(initial_model, batch, learning_rate):
# Define a group of model variables and set them to `initial_model`. Must
# be defined outside the @tf.function.
model_vars = collections.OrderedDict([
(name, tf.Variable(name=name, initial_value=value))
for name, value in initial_model.items()
])
optimizer = tf.keras.optimizers.SGD(learning_rate)
@tf.function
def _train_on_batch(model_vars, batch):
# Perform one step of gradient descent using loss from `batch_loss`.
with tf.GradientTape() as tape:
loss = forward_pass(model_vars, batch)
grads = tape.gradient(loss, model_vars)
optimizer.apply_gradients(
zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))
return model_vars
return _train_on_batch(model_vars, batch)
print(str(batch_train.type_signature))
输出：
(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>,float32> -> <weights=float32[784,10],bias=float32[10]>)
最先看到，输入参数的类型是有显式声明的。
函数体中先声明了变量model_vars并赋值，相当于对initial_model做了次深拷贝。再实例化一个SGD优化器optimizer。在_train_on_batch函数中，我们调用了之前定义的前向传播函数forward_pass来计算loss，然后显式地计算反向传播的梯度，再显式地apply梯度到模型变量上，最终返回更新后的模型参数。
简单来说，batch_train就是输入当前模型、一个batch的数据和学习率，返回一组优化后的模型参数。
接下来我们用上面定义的初始模型（参数全为0），随便拿一个batch出来训练一下：
model = initial_model
losses = []
for _ in range(5):
model = batch_train(model, sample_batch, 0.1)
losses.append(batch_loss(model, sample_batch))
print("5 loops loss:", losses)
上面的代码中，我们用同一份batchsample_batch去训练一个初始模型。可以发现，每次调用batch_train的过程，就是模型迭代一轮的过程，下一轮输入的模型就是上一轮返回的模型。正常执行的话，应该会看到类似下面的输出：
5 loops loss: [0.19690022, 0.13176315, 0.10113226, 0.08273813, 0.07030139]
可以看出loss值是在不断下降的。而且重复跑多少次都是类似的结果，因为batch_train函数中有做深拷贝，不会影响函数外的实参。
local训练
接下来我们写一个函数，能够对一系列的batch进行学习，并且更新模型：
LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)
@tff.federated_computation(MODEL_TYPE, tf.float32, LOCAL_DATA_TYPE)
def local_train(initial_model, learning_rate, all_batches):
# Mapping function to apply to each batch.
@tff.federated_computation(MODEL_TYPE, BATCH_TYPE)
def batch_fn(model, batch):
return batch_train(model, batch, learning_rate)
return tff.sequence_reduce(all_batches, initial_model, batch_fn)
print(str(local_train.type_signature))
输出：
(<<weights=float32[784,10],bias=float32[10]>,float32,<x=float32[?,784],y=int32[?]>*> -> <weights=float32[784,10],bias=float32[10]>)
这个函数实现了对当前模型initial_model进行一系列batch的更新。可以看到，local_train输入的参数都不是联邦类型的，但是它却用federated_computation来包装，一是为了表达一种“联邦”的逻辑概念，二是如果内外两个函数都改成tf_computation后也会报错，因为调用了TFF框架的函数tff.sequence_reduce()，不是TF的代码。
训练一次试试：
locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5])
这里我们喂给初始模型所有的数字5的训练集，但没有其他数字，我们待会看看这种极度不平衡的训练集分布（non-IID），会带来什么影响。
定义评价函数
我们想知道模型究竟训练得咋样了，我们需要定义一个评价函数，来量化当前模型对数据样本的分类表现：
@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)
def local_eval(model, all_batches):
# TODO(b/120157713): Replace with `tff.sequence_average()` once implemented.
return tff.sequence_sum(
tff.sequence_map(
tff.federated_computation(lambda b: batch_loss(model, b), BATCH_TYPE),
all_batches))
print(str(local_eval.type_signature))
输出函数定义类型：
(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>*> -> float32)
函数local_eval对当前模型在所有batch上的loss进行求和计算，我们关注一下tff.federated_computation的函数型调用方式，可以定义一个匿名的联邦计算函数。
对于初始化模型，以及上面训练了一轮的locally_trained_model，我们计算输出它的loss值和评价指标：
print('initial_model loss [num 5] =', local_eval(initial_model, federated_train_data[5]))
print('locally_trained_model loss [num 5] =', local_eval(locally_trained_model, federated_train_data[5]))
print('initial_model loss [num 0] =', local_eval(initial_model, federated_train_data[0]))
print('locally_trained_model loss [num 0] =', local_eval(locally_trained_model, federated_train_data[0]))
上面的代码正常执行的话，可以看到如下结果：
initial_model loss [num 5] = 23.025854
locally_trained_model loss [num 5] = 0.4348469
initial_model loss [num 0] = 23.025854
locally_trained_model loss [num 0] = 74.500755
通过结果我们发现，我们的模型locally_trained_model因为在数字5的数据集上进行训练，因此对数字5的识别程度较高，loss降低到了0.4348469。但是因为没有见过数字0，识别度较低，loss反而比初始模型增大了。
定义联邦评价函数
为了解决上述问题，我们开始正式的联邦训练：用10个clients分别学习10个数字的样本，得到一个全方面的模型。首先定义联邦学习的模型的评价函数：
SERVER_MODEL_TYPE = tff.FederatedType(MODEL_TYPE, tff.SERVER)
CLIENT_DATA_TYPE = tff.FederatedType(LOCAL_DATA_TYPE, tff.CLIENTS)
@tff.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)
def federated_eval(model, data):
return tff.federated_mean(
tff.federated_map(local_eval, [tff.federated_broadcast(model), data]))
print(str(federated_eval.type_signature))
输出：
(<<weights=float32[784,10],bias=float32[10]>@SERVER,{<x=float32[?,784],y=int32[?]>*}@CLIENTS> -> float32@SERVER)
到了联邦的场景，终于需要显式定义联邦类型了。这里我们定义了两种类型：SERVER_MODEL_TYPE是存放在S端的模型参数，是全局唯一的。CLIENT_DATA_TYPE是存在C端的，默认全局不唯一。
函数federated_eval把模型丢给clients去计算一下loss。这里体现了联邦学习的一个关键点：在现实场景中，data是存放在C端且很隐私的，因此用data去评价模型只能给C端去做。那么，tff.federated_broadcast函数就负责完成这个S端到C端的传输过程，让我们无需关注网络传输等细节。
对于初始模型和上面训练了一轮的模型，我们可以跑个全局评价看看：
print('initial_model loss =', federated_eval(initial_model,
federated_train_data))
print('locally_trained_model loss =',
federated_eval(locally_trained_model, federated_train_data))
输出：
initial_model loss = 23.025852
locally_trained_model loss = 54.43263
可以看出，单个节点训练的效果还不如全0初始化。因此，如何整合多方训练，是联邦学习能达到好效果的关键。
定义联邦训练
对于联邦训练，我们定义如下训练函数：
SERVER_FLOAT_TYPE = tff.FederatedType(tf.float32, tff.SERVER)
@tff.federated_computation(SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE,
CLIENT_DATA_TYPE)
def federated_train(model, learning_rate, data):
return tff.federated_mean(
tff.federated_map(local_train, [
tff.federated_broadcast(model),
tff.federated_broadcast(learning_rate), data
]))
其中，类型SERVER_FLOAT_TYPE的解读交给读者完成。federated_train函数其实是把S端的模型和学习率通过tff.federated_broadcast“转化”为C端的数据，再分别放置到local_train里进行训练，得到10个更新后的模型后，用tff.federated_mean直接求平均值作为全局模型的更新。此时，data所代表的10类数字样本分别给予10个clients去训练，即体现了联邦学习应对non-IID数据的场景。
最后用一段循环，来对S端模型进行迭代训练：
model = initial_model
learning_rate = 0.1
for round_num in range(5):
# 每一轮，把大家的模型分别更新一下，取平均之后拿回来（做赋值替换）
model = federated_train(model, learning_rate, federated_train_data)
# 把学习率减小一点
learning_rate = learning_rate * 0.9
# 算个loss输出一下
loss = federated_eval(model, federated_train_data)
print('round {}, loss={}'.format(round_num, loss))
# 下一轮S端的模型又发给各位clients去更新
可以得到类似的输出：
round 0, loss=21.60552406311035
round 1, loss=20.365678787231445
round 2, loss=19.27480125427246
round 3, loss=18.31110954284668
round 4, loss=17.45725440979004
模型迭代训练的流程已经写到代码注释里了。看到了这里，请各位同学结合文章开头的那张图片想象一下这个联邦训练过程，如果能把代码和图片每个部分对应上，那么应该就算入门了。
5、后记
如果不是英语苦手的同学请一定去看一下谷歌官方的文档和教程~。接下来我的学习与博客计划包括：
- 高级TFF API的学习、练习
- 用低级API实现复杂的联邦学习模型
感谢你的阅读！若有疑问或建议请在评论区交流~
2020/3/31 更新：由于旧版TFF惨遭淘汰，因此本文参考最新文档，重新整理了代码。现在应该可以正常运行了，请各位读者重新安装最新版（v0.13.1）的TFF库，谢谢各位的留言！
友情链接：楼楼表情包库
6、完整代码
上面用作讲解的例子的代码，参考自谷歌TFF文档。
import collections
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
tf.compat.v1.enable_v2_behavior()
# TODO(b/148678573,b/148685415): must use the ReferenceExecutor because it
# supports unbounded references and tff.sequence_* intrinsics.
tff.framework.set_default_executor(tff.framework.ReferenceExecutor())
mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
NUM_EXAMPLES_PER_USER = 1000
BATCH_SIZE = 100
def get_data_for_digit(source, digit):
output_sequence = []
all_samples = [i for i, d in enumerate(source[1]) if d == digit]
for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):
batch_samples = all_samples[i:i + BATCH_SIZE]
output_sequence.append({
'x':
np.array([source[0][i].flatten() / 255.0 for i in batch_samples],
dtype=np.float32),
'y':
np.array([source[1][i] for i in batch_samples], dtype=np.int32)
})
return output_sequence
federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)]
federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]
BATCH_SPEC = collections.OrderedDict(
x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),
y=tf.TensorSpec(shape=[None], dtype=tf.int32))
BATCH_TYPE = tff.to_type(BATCH_SPEC)
print(str(BATCH_TYPE))
MODEL_SPEC = collections.OrderedDict(
weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),
bias=tf.TensorSpec(shape=[10], dtype=tf.float32))
MODEL_TYPE = tff.to_type(MODEL_SPEC)
print(MODEL_TYPE)
@tf.function
def forward_pass(model, batch):
predicted_y = tf.nn.softmax(
tf.matmul(batch['x'], model['weights']) + model['bias'])
return -tf.reduce_mean(
tf.reduce_sum(
tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE)
def batch_loss(model, batch):
return forward_pass(model, batch)
initial_model = collections.OrderedDict(
weights=np.zeros([784, 10], dtype=np.float32),
bias=np.zeros([10], dtype=np.float32))
sample_batch = federated_train_data[5][-1]
print(batch_loss(initial_model, sample_batch))
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def batch_train(initial_model, batch, learning_rate):
# Define a group of model variables and set them to `initial_model`. Must
# be defined outside the @tf.function.
model_vars = collections.OrderedDict([
(name, tf.Variable(name=name, initial_value=value))
for name, value in initial_model.items()
])
optimizer = tf.keras.optimizers.SGD(learning_rate)
@tf.function
def _train_on_batch(model_vars, batch):
# Perform one step of gradient descent using loss from `batch_loss`.
with tf.GradientTape() as tape:
loss = forward_pass(model_vars, batch)
grads = tape.gradient(loss, model_vars)
optimizer.apply_gradients(
zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))
return model_vars
return _train_on_batch(model_vars, batch)
print(str(batch_train.type_signature))
model = initial_model
losses = []
for _ in range(5):
model = batch_train(model, sample_batch, 0.1)
losses.append(batch_loss(model, sample_batch))
print("5 loops loss:", losses)
LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)
@tff.federated_computation(MODEL_TYPE, tf.float32, LOCAL_DATA_TYPE)
def local_train(initial_model, learning_rate, all_batches):
# Mapping function to apply to each batch.
@tff.federated_computation(MODEL_TYPE, BATCH_TYPE)
def batch_fn(model, batch):
return batch_train(model, batch, learning_rate)
return tff.sequence_reduce(all_batches, initial_model, batch_fn)
print(str(local_train.type_signature))
locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5])
@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)
def local_eval(model, all_batches):
# TODO(b/120157713): Replace with `tff.sequence_average()` once implemented.
return tff.sequence_sum(
tff.sequence_map(
tff.federated_computation(lambda b: batch_loss(model, b), BATCH_TYPE),
all_batches))
print(str(local_eval.type_signature))
print('initial_model loss [num 5] =', local_eval(initial_model, federated_train_data[5]))
print('locally_trained_model loss [num 5] =', local_eval(locally_trained_model, federated_train_data[5]))
print('initial_model loss [num 0] =', local_eval(initial_model, federated_train_data[0]))
print('locally_trained_model loss [num 0] =', local_eval(locally_trained_model, federated_train_data[0]))
SERVER_MODEL_TYPE = tff.FederatedType(MODEL_TYPE, tff.SERVER)
CLIENT_DATA_TYPE = tff.FederatedType(LOCAL_DATA_TYPE, tff.CLIENTS)
@tff.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)
def federated_eval(model, data):
return tff.federated_mean(
tff.federated_map(local_eval, [tff.federated_broadcast(model), data]))
print(str(federated_eval.type_signature))
print('initial_model loss =', federated_eval(initial_model,
federated_train_data))
print('locally_trained_model loss =',
federated_eval(locally_trained_model, federated_train_data))
SERVER_FLOAT_TYPE = tff.FederatedType(tf.float32, tff.SERVER)
@tff.federated_computation(SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE,
CLIENT_DATA_TYPE)
def federated_train(model, learning_rate, data):
return tff.federated_mean(
tff.federated_map(local_train, [
tff.federated_broadcast(model),
tff.federated_broadcast(learning_rate), data
]))
model = initial_model
learning_rate = 0.1
for round_num in range(5):
# 每一轮，把大家的模型分别更新一下，取平均之后拿回来（做赋值替换）
model = federated_train(model, learning_rate, federated_train_data)
# 把学习率减小一点
learning_rate = learning_rate * 0.9
# 算个loss输出一下
loss = federated_eval(model, federated_train_data)
print('round {}, loss={}'.format(round_num, loss))
# 下一轮S端的模型又发给各位clients去更新
扫一扫
4037
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
