Tensorflow Federated Framework 谷歌联邦学习框架:自底向上简明入门

0、前言

参考:谷歌联邦学习框架谷歌联邦学习BlogStack Overflow上的讨论

只使用TF实现联邦学习的另两篇博客:【联邦学习】用Tensorflow实现联邦模型AlexNet on CIFAR-10【联邦学习】用Tensorflow实现联邦学习传递梯度

计划

目前中文社区在联邦学习这一块大多还是科普性的文章,本文希望能从代码实现方面来介绍联邦学习,也同时作为我个人的一份学习笔记,如有描述错误的地方请多多指正,共同学习!

联邦学习概念

联邦学习是谷歌在2016年提出的概念:在分布式的场景下,训练数据分别保存在每个clients中,希望提出一种训练方法:跨多个参与客户端(clients)训练一个共享的全局模型。其中的重点关注的问题包括:

  1. 参与的clients数量众多,比如每个人的手机、物联网设备等
  2. clients不能一直保持在线状态,比如只有当手机接入充电器和wifi时才会参与计算
  3. 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框架中,也有两个不同层次的接口:

  1. Federated Learning (FL) API:该层提供了一组高阶接口,使开发者能够将包含的联合训练和评估实现应用于现有的 TensorFlow 模型。
  2. 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接受三种类型的输入:listtuplecollections.OrderedDict(也就是collection.namedtuple生成的subclass生成的对象,一种有序的字典dict类型)
将会输出:

<int8,int16>
<key1=float32,int8,key3=string*>
<y=float32,x=string[3,4,5]>

其中的尖括号是named tuple types的标志。可以注意到,元素的key可选的,有keykey的元素位置也是很随意的。一般来说,这个类可以用于定义模型的参数、输入输出等。

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。它指明了输入输出的形式,用圆括号和箭头->表示。同时,只有一个输入和一个输出,但是输入变量可以是复合类型(例如这里的列表类型),因此是能满足各种需求的。

端有关的类型

端有关的类型,是联邦学习逻辑层面的重点。它们主要完成两件任务:

  1. 显式地定义数据值应该存放在C端还是S端(Placement
  2. 定义这个数据是否全局一致(All equal?

我们可以把端有关的类型,想象成一个快递盒:里面装的数据类型是上文提到的端无关类型,盒子上贴着两个标签PlacementAll equal?。那么联邦学习的流程,可以类比成:快递盒子从S端中心节点发出,交给C端加工厂处理(盒子由S端传输到C端,还可以复制成多份)。加工厂把盒子里的东西拿出来操作(与端无关)。操作完成后再打包发回中心节点(盒子由C端传输到S端),中心点把各个工厂发回的加工品聚合成一个,就完成了一轮加工。

Placement type

顾名思义,就是定义存放位置的类型。目前,我们能用到的有两个:tff.SERVERtff.CLIENTS,也就是定义了S端和C端,把他们俩当常数用就好了。

以后也许可以定义新的Placement,以实现更复杂的场景。具体咋用呢,看下面的联邦类型。

Federated types(tff.FederatedType

以数据驱动的联邦学习,终于到了定义联邦类型的时候了。联邦类型tff.FederatedType把上面提到的端无关类型包装起来,并增加两个属性:

  1. placement(必填),必须是tff.SERVERtff.CLIENTS这些Placement type
  2. all_equal(可选,默认为None)。类型为bool,代表着这份数据是否全局统一,还是可以有不同的值。如果没有指定all_equal,它会根据placement的值来选择。默认情况下placement=tff.SERVERall_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

注意到两个地方:

  1. 返回值类型是自动推算出来的(int32
  2. 变量的自动类型转换是可以的(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)

我们可以这样理解:

  1. tff.tf_computation里定义的是纯TF逻辑
  2. tff.federated_computation里可以是TF逻辑(可以调用tff.tf_computation定义的函数),也可以是联邦训练的逻辑(调用其他端有关函数,包括很多TFF内置函数)

内置函数

TFF框架还提供了许多内置函数,例如上面用到的tff.federated_mean。它们也可以分为端有关和端无关的,区分的方法非常简单,如果函数是tff.federated_开头的,就是端有关的,会规定输入输出的存放位置。其他的函数即是端无关的,输入输出的类型也是端无关类型(有些也可以用在端有关场景)。具体有哪些函数可以用,请查阅官方文档:https://www.tensorflow.org/federated/api_docs/python/tff#functions
官方文档函数列表

其他注意

  1. 函数调用的语法就和调用普通的python函数一样,但是要注意输入参数的类型是否匹配(或可以自动转换)
  2. tff.tf_computation包装的函数不能调用tff.federated_computation包装的,但反过来可以
  3. 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分类数据集。可视化如下图:
mnist样本
上面预处理的时候,做了几件事情:

  1. 把每个样本从矩阵变成一维向量(长度为28x28=784),元素从灰度(0~255)缩放到0~1的范围
  2. 把同一个数字的样本手动挑出来放到一个list里,并且分了batch
  3. 得到的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_TYPEMODEL_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、后记

如果不是英语苦手的同学请一定去看一下谷歌官方的文档和教程~。接下来我的学习与博客计划包括:

  1. 高级TFF API的学习、练习
  2. 用低级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去更新
评论 107
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值