RSVP for your your local TensorFlow Everywhere event today!

自定义联合算法,第 1 部分:Federated Core 简介

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 中查看源代码

本系列教程包括两个部分,此为第一部分。该系列演示了如何使用 Federated Core (FC) 在 TensorFlow Federated (TFF) 中实现自定义类型的联合算法。Federated Core 是一组较低级别的接口,这些接口是我们实现联合学习 (FL) 层的基础。

第一部分更具概念性;我们将介绍 TFF 中使用的一些关键概念和编程抽象,并在一个非常简单的示例(分布式温度传感器阵列)中演示它们的用法。在本系列的第二部分中,我们将使用此处介绍的机制来实现一个简单版本的联合训练和评估算法。我们鼓励您稍后在 tff.learning 中研究联合平均的实现

在本系列的最后,您应该能够认识到 Federated Core 的应用并不仅限于学习。我们提供的编程抽象非常通用,例如可用于对分布式数据进行分析和其他自定义类型的计算。

尽管本教程可独立使用,但我们建议您先阅读有关图像分类文本生成的教程,获得对 TensorFlow Federated 框架和 Federated Learning API (tff.learning) 更高级和更循序渐进的介绍,它将帮助您在上下文中理解我们在此介绍的概念。

预期用途

简而言之,Federated Core (FC) 是一种开发环境,可以紧凑地表达将 TensorFlow 代码与分布式通信算子(比如在联合平均中使用的算子)相结合的程序逻辑。它可以在系统中的一组客户端设备上计算分布式总和、平均值和其他类型的分布式聚合,向这些设备广播模型和参数等。

您可能知道 tf.contrib.distribute,此时自然会问的一个问题可能是:该框架在哪些方面有所不同?毕竟,两种框架都试图使 TensorFlow 进行分布式计算。

其中一种思路是,tf.contrib.distribute 的既定目标是允许用户以最小的更改使用现有模型和训练代码实现分布式训练,且大部分重点放在如何利用分布式架构来提高现有训练代码的效率。TFF 的 Federated Core 的目标是使研究员和从业者能够明确控制将在系统中使用的分布式通信的具体模式。FC 的重点是提供一种灵活可扩展的语言来表达分布式数据流算法,而不是具体的一组已实现的分布式训练能力。

TFF 的 FC API 的主要目标受众之一是研究员和从业者,他们可能希望尝试新的联合学习算法,并评估微妙的设计选择(这些选择会影响分布式系统中数据流的编排方式)的结果,但又不想被系统实现细节所困扰。FC API 所针对的抽象级别大致对应于伪代码,可以用来描述研究论文中的联合学习算法的机制(系统中存在什么数据以及如何对其进行转换),但又不降低到单个点对点网络消息交换的级别。

TFF 总体上针对的是数据分布的场景(并且出于隐私等原因必须保持这种状态),以及可能无法在某个集中位置收集所有数据的场景。与所有数据都可以在数据中心积累到某个集中位置的场景相比,这意味着机器学习算法的实现需要提高显式控制的程度。

准备工作

在深入研究代码之前,请尝试运行以下 “Hello World” 示例,以确保您的环境已正确设置。如果无法正常运行,请参阅安装指南中的说明。

!pip install --quiet --upgrade tensorflow_federated_nightly
!pip install --quiet --upgrade nest_asyncio

import nest_asyncio
nest_asyncio.apply()
import collections

import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
@tff.federated_computation
def hello_world():
  return 'Hello, World!'

hello_world()
b'Hello, World!'

联合数据

TFF 的一个显著功能是,您可以通过它在联合数据上紧凑地表达基于 TensorFlow 的计算。在本教程中,我们将使用联合数据这一术语来指代托管在分布式系统中一组设备上的数据项的集合。例如,在移动设备上运行的应用可以收集数据并将其存储在本地,而无需上传到某一集中位置。或者,分布式传感器阵列可以在本地收集并存储温度读数。

像上面示例中这样的联合数据在 TFF 中被视为一等公民,即,它们可以显示为函数的参数和结果,并具有类型。为了强化这一概念,我们将联合数据集称为联合值联合类型的值

需要理解的一个要点是,我们将所有设备上的数据项的整个集合(例如,来自分布式阵列中所有传感器温度读数的整个集合)建模为单个联合值。

例如,下面是在 TFF 中定义由一组客户端设备托管的联合浮点类型的方式。可以将分布式传感器阵列中出现的温度读数的集合建模为此联合类型的值。

federated_float_on_clients = tff.FederatedType(tf.float32, tff.CLIENTS)

更普遍的是,TFF 中的联合类型是通过指定其成员组成(留驻在各个设备上的数据项)的类型 T 和托管此类型联合值的设备组 G(加上我们会在稍后提及的第三个可选位)来定义的。我们将托管联合值的设备组 G 称为该值的布局。因此,tff.CLIENTS 是布局示例。

str(federated_float_on_clients.member)
'float32'
str(federated_float_on_clients.placement)
'CLIENTS'

带有成员组成 T 和布局 G 的联合类型可以紧凑地表示为 {T}@G,如下所示。

str(federated_float_on_clients)
'{float32}@CLIENTS'

此简明表示法中的大括号 {} 提醒您成员组成(不同设备上的数据项)可能有所不同,例如您所期望的温度传感器读数,因此,客户端会作为整体共同托管 T 型项的多重集,它们共同构成联合值。

需要注意的是,联合值的成员组成通常对程序员不透明,即不应将联合值视为由系统中的设备标识符进行键控的简单 dict,这些值应仅由抽象表示各种分布式通信协议(如聚合)的联合算子进行集体转换。如果这听起来太过抽象,不要担心,我们稍后将回到这个问题,并用具体示例对其进行演示。

TFF 中的联合类型有两种形式:一种是联合值的成员组成可能不同(如上所示),另一种是联合值的成员组成全部相等。这由 tff.FederatedType 构造函数中的第三个可选参数 all_equal(默认为 False)来控制。

federated_float_on_clients.all_equal
False

可以将带有布局 G,且其中所有 T 型成员组成已知相等的联合类型紧凑地表示为 T@G(与 {T}@G 相对,也就是说,去掉大括号表示成员组成的多重集由单个项目构成)。

str(tff.FederatedType(tf.float32, tff.CLIENTS, all_equal=True))
'float32@CLIENTS'

在实际场景中可能会出现的此类联合值的一个示例是超参数(例如学习率、裁剪范数等),该超参数已由服务器广播到参与联合训练的一组设备。

另一个示例是在服务器上预先训练的一组机器学习模型参数,然后将其广播到一组客户端设备,它们可以在这组设备上针对每个用户进行个性化设置。

例如,假设对于一个简单的一维线性回归模型,我们有一对 float32 参数 ab。我们可以构造如下用于 TFF 的(非联合)类型的此类模型。打印的类型字符串中的尖括号 <> 是命名或未命名元组的紧凑 TFF 表示法。

simple_regression_model_type = (
    tff.StructType([('a', tf.float32), ('b', tf.float32)]))

str(simple_regression_model_type)
'<a=float32,b=float32>'

请注意,虽然我们在上文中仅指定了 dtype,但也支持非标量类型。在上面的代码中,tf.float32 是更通用的 tff.TensorType(dtype=tf.float32, shape=[]) 的快捷表示法。

将此模型广播到客户端时,生成的联合值类型可以用如下方法表示。

str(tff.FederatedType(
    simple_regression_model_type, tff.CLIENTS, all_equal=True))
'<a=float32,b=float32>@CLIENTS'

根据上面的联合浮点的对称性,我们将这种类型称为联合元组。一般来说,我们会经常使用术语联合 XYZ 来指代联合值,其中成员组成类似 XYZ。因此,我们将对联合元组联合序列联合模型等进行讨论。

现在,回到 float32@CLIENTS,尽管它看起来是在多个设备上复制的,但它实际上是单一的 float32,因为所有成员都相同。通常,您可能会想到任意全等联合类型(即一种 T@G 形式)与非联合类型 T 同构,因为这两种情况实际上都只有 T 类型的单个(尽管可能是复制的)项。

鉴于 TT@G 之间的同构性,您可能想知道后一种类型能够起到什么作用(如果有的话)。请继续阅读。

布局

设计概述

在上一部分中,我们介绍了布局的概念(可能会共同托管联合值的系统参与者组),并且我们还演示了将 tff.CLIENTS 用作布局的示例规范。

为了解释为什么放置的概念如此重要,以至于我们需要将其合并到 TFF 类型系统中,请回想一下本教程开始时提到的有关 TFF 某些预期用途的内容。

尽管在本教程中,您只会看到在模拟环境中本地执行的 TFF 代码,但我们的目标是使 TFF 能够编写可部署在分布式系统中的物理设备组(可能包括运行 Android 的移动或嵌入式设备)上执行的代码。其中,每个设备都将接收单独的一组指令以在本地执行,具体取决于它在系统中所扮演的角色(最终用户设备、集中协调器、多层架构中的中间层等)。重要的是要能够推断出哪些设备子集执行什么代码,以及数据的不同部分可能在哪里进行物理实现。

当处理移动设备上的应用数据时,这一点尤其重要。由于数据私有且可能敏感,我们要能静态验证此类数据永远不会离开设备(并证明对数据进行处理的实际方式)。布局规范是为此目的而设计的一种机制。

TFF 是一种以数据为中心的编程环境,正因为如此,它与一些现有框架不同,这些框架专注于运算和这些运算可能运行的位置,而 TFF 专注于数据、数据实现的位置,以及转换方式。因此,布局被建模为 TFF 中数据的属性,而不是数据上运算的属性。的确,您将在下一部分中看到,某些 TFF 运算会跨位置,并且可以说是“在网络中”运行,而不是由一台或一组机器执行。

将某个值的类型表示为 T@G{T}@G(而不仅仅是 T)可使数据布局决策明确,并且搭配 TFF 中编写的程序的静态分析,它可以作为为设备端敏感数据提供形式上的隐私保障的基础。

但此时需要注意,虽然我们鼓励 TFF 用户明确托管数据的参与设备(布局),但程序员永远不会处理原始数据或各个参与者的身份。

注:虽然超出了本教程的范围,但我们应该提到,以上内容有一个明显的例外,即 tff.federated_collect 算子(它旨在用作低级基元,并仅用于特殊情况)。不建议在可以避免的情况下显式使用它,因为它可能会限制将来可能的应用。例如,如果在静态分析过程中,我们决定让某个计算使用这种低级机制,我们可能会禁止其访问某些类型的数据。

在 TFF 代码的主体内,根据设计,无法枚举构成由 tff.CLIENTS 表示的组的设备,也无法探测组中是否存在某个特定设备。在 Federated Core API、基础架构抽象集或我们提供的用于支持模拟的核心运行时基础结构中,没有任何设备或客户端标识的概念。您编写的所有计算逻辑都将表达为在整个客户端组上的运算。

回想一下我们前面提到的联合类型的值与 Python dict 的不同,因为它不能简单地枚举其成员组成。将您的 TFF 程序逻辑所操作的值视为与布局(组)所关联,而不是与单个参与者相关联。

布局在 TFF 中也设计为一等公民,并且可以显示为 placement 类型的参数和结果(通过 API 中的 tff.PlacementType 表示)。将来,我们计划提供各种算子来转换或合并布局,但这不在本教程的讨论范围之内。就目前而言,将 placement 视为 TFF 中的不透明基元内置类型就足够了(类似于 Python 中的不透明内置类型 intbool),其中 tff.CLIENTS 是该类型的常量文字, 与 1int 类型的常量文字没什么区别。

指定布局

TFF 提供了两种基本的布局文本 tff.CLIENTStff.SERVER,使用通过单个集中式服务器协调器编排的多种客户端设备(移动电话、嵌入式设备、分布式数据库、传感器等),使自然建模为客户端-服务器架构的各种实际场景易于表达。TFF 的设计还支持自定义位置、多客户端组、多层和其他更通用的分布式架构,但对这些内容的讨论不在本教程的范围之内。

TFF 没有规定 tff.CLIENTStff.SERVER 实际代表的内容。

特别是,tff.SERVER 可能是单个物理设备(单一实例组的成员),但它也可能是运行状态机复制的容错集群中的一组副本,我们不做任何特殊架构的假设。相反,我们使用前一部分提到的 all_equal 位来表达我们通常只在服务器上处理单个数据项这一事实。

同样,在某些应用中,tff.CLIENTS 可能代表系统中的所有客户端,在联合学习的上下文中,我们有时将其称为群体,但在联合平均的生产实现这个示例中,它可能代表队列(选择参加某轮训练的客户端的子集)。当部署计算以执行(或者就像模型环境中的 Python 函数那样被调用)时,其中的抽象定义布局将被赋予具体含义。在我们的本地模拟中,客户端组由作为输入提供的联合数据来确定。

联合计算

声明联合计算

TFF 是支持模块化开发的强类型函数式编程环境。

TFF 中的基本组成单位是联合计算,它是可以接受联合值作为输入并返回联合值作为输出的逻辑的一部分。下面的代码定义了一个计算,它计算的是前一个示例中传感器阵列报告的平均温度。

@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(sensor_readings):
  return tff.federated_mean(sensor_readings)

查看上面的代码,此时您可能会问:TensorFlow 中不是已经有用于定义可组合单元的装饰器构造(如 tf.function)了吗?既然如此,为什么还要引入另一个构造?它们有什么区别?

简短回答是,tff.federated_computation 封装容器生成的代码既不是 TensorFlow,也不是 Python,它是一种与平台无关的内部胶水语言中的分布式系统规范。这听起来确实很神秘,但请牢记这一将联合计算作为分布式系统抽象规范的直观解释。我们稍后将对其进行说明。

首先,我们来思考一下定义。TFF 计算通常会被建模为函数,参数可有可无,但要有明确定义的类型签名。您可以通过查询计算的 type_signature 属性来打印计算的类型签名,如下所示。

str(get_average_temperature.type_signature)
'({float32}@CLIENTS -> float32@SERVER)'

类型签名告诉我们,该计算接受客户端设备上不同传感器读数的集合,并在服务器上返回单个平均值。

在进一步讨论之前,让我们先思考一个问题:此计算的输入和输出位于不同位置(在 CLIENTS 上和在 SERVER 上)。回想一下我们在上一部分所讲的关于 TFF 运算如何跨位置并在网络中运行的内容,以及我们刚才讲到的有关联合计算表示分布式系统抽象规范的内容。我们刚刚定义了这样一种计算:一个简单的分布式系统,其中数据在客户端设备上使用,而聚合结果出现在服务器上。

在许多实际场景中,代表顶级任务的计算倾向于接受其输入并在服务器上报告其输出,这表明,计算可能会由在服务器上发起和终止的查询所触发。

但是,FC API 并不强制实施此假设,并且我们在内部使用的许多构建块(包括您可能在 API 中见到的许多 tff.federated_... 算子)的输入和输出都有不同的布局,因此,通常来说,您不应将联合计算视为在服务器上运行由服务器执行的内容。服务器只是联合计算中的一种类型的参与者。在思考此类计算的机制时,最好始终默认使用全局网络范围的视角,而不是使用单个集中协调器的视角。

一般来说,对于输入和输出的类型 TU,函数类型签名会分别被紧凑地表示为 (T -> U)。形式参数的类型(此处为 sensor_readings)被指定为装饰器的参数。您无需指定结果的类型,因为它会自动确定。

尽管 TFF 确实提供了有限形式的多态性,但我们强烈建议程序员明确自己使用的数据类型,因为这样可以更轻松地理解、调试和在形式上验证您的代码的属性。在某些情况下,必须明确指定类型(例如,当前无法直接执行多态计算时)。

执行联合计算

为了支持开发和调试,TFF 允许您直接调用以此方式定义为 Python 函数的计算,如下所示。对于 all_equal 位设置为 False 的情况,您可以将其作为 Python 中的普通 list 进行馈送,而对于 all_equal 位设置为 True 的情况,您可以直接馈送(单)成员组成。这也是向您反馈结果的方式。

get_average_temperature([68.5, 70.3, 69.8])
69.53334

在模拟模式下运行此类计算时,您将充当具有系统范围视图的外部观察者,您能够在网络中的任何位置提供输入和使用输出,这里正是如此,您提供了客户端值作为输入,并使用了服务器结果。

现在,让我们回到先前关于 tff.federated_computation 装饰器用胶水语言发出代码的注释。尽管 TFF 计算的逻辑可以用 Python 表达为普通函数(您只需按照上述方法,使用 tff.federated_computation 对其进行装饰),而且您可以像此笔记本中的其他 Python 函数一样直接调用它们,但在后台,正如我们前面提到的,TFF 计算实际上不是 Python。

我们的意思是,当 Python 解释器遇到一个用 tff.federated_computation 装饰的函数时,它会对此函数主体中的语句进行一次跟踪(在定义时),然后构造该计算逻辑的序列化表示以供将来使用(无论是用于执行,还是作为子组件合并到另一个计算中)。

您可以通过添加打印语句来验证这一点,如下所示:

@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(sensor_readings):

  print ('Getting traced, the argument is "{}".'.format(
      type(sensor_readings).__name__))

  return tff.federated_mean(sensor_readings)
Getting traced, the argument is "ValueImpl".

您可以将定义了联合计算的 Python 代码想象成在非 Eager 上下文中构建了 TensorFlow 计算图的 Python 代码(如果您不熟悉 TensorFlow 的非 Eager 用法,请想象您的 Python 代码定义了运算的计算图以稍后执行,但实际上并不立即运行它们)。TensorFlow 中的非 Eager 计算图构建代码是 Python,但用此代码构建的 TensorFlow 计算图与平台无关且可序列化。

同样,TFF 计算用 Python 进行定义,但其主体中的 Python 语句(如我们刚才展示的示例中的 tff.federated_mean)会在后台被编译成可移植、与平台无关和可序列化的表示形式。

作为开发者,您无需关注此表示形式的细节,因为您不会直接使用它,但您应该知道它的存在,以及 TFF 计算在本质上非 Eager,且无法捕获任意 Python 状态。TFF 计算主体中包含的 Python 代码会在定义时(即用 tff.federated_computation 装饰的 Python 函数在序列化之前被跟踪时)执行。调用时不会再次对其进行跟踪(函数为多态时除外;有关详细信息,请参阅文档页面)。

您可能想知道为什么我们选择引入专用的内部非 Python 表示形式。其中一个原因是,TFF 计算的最终目的是可部署到实际物理环境中,并托管在可能无法使用 Python 的移动或嵌入式设备上。

另一个原因是 TFF 计算表达的是分布式系统的全局行为,而 Python 程序表达的是各个参与者的本地行为。您可以在上面的简单示例中看到这一点,即使用特殊算子 tff.federated_mean 接受客户端设备上的数据,但将结果存储在服务器上。

无法将算子 tff.federated_mean 轻松建模为 Python 中的普通算子,因为它不在本地执行,如前所述,它表示协调多个系统参与者行为的分布式系统。我们将此类算子称为联合算子,以将其与 Python 中的普通(本地)算子进行区分。

因此,TFF 类型系统以及 TFF 语言支持的基本运算集与 Python 中的大不相同,因此必须使用专用表示形式。

组成联合计算

如上所述,最好将联合计算及其组成理解为分布式系统的模型,并且可以将联合计算的组成过程想象成从较简单的分布式系统组成较复杂的分布式系统的过程。您可以将 tff.federated_mean 算子视为一种具有类型签名 ({T}@CLIENTS -> T@SERVER) 的内置模板联合计算(实际上,就像您编写的计算一样,该算子的结构也很复杂,我们会在后台把它分解成更简单的算子)。

组成联合计算的过程也是如此。可以在另一个用 tff.federated_computation 装饰的 Python 函数主体中调用计算 get_average_temperature,这样做会将其嵌入到父级的主体中,这与先前将 tff.federated_mean 嵌入到其自身主体中的方式相同。

需要注意的一个重要限制是,用 tff.federated_computation 装饰的 Python 函数的主体必须由联合算子组成(即它们不能直接包含 TensorFlow 运算)。例如,不能直接使用 tf.nest 接口添加一对联合值。TensorFlow 代码仅限用 tff.tf_computation 装饰的代码块,下一部分将对此进行讨论。只有以这种方式封装,才能在 tff.federated_computation 主体中调用封装后的 TensorFlow 代码。

这样区分是出于技术原因(很难欺骗 tf.add 之类的算子来使用非张量),以及架构原因。联合计算的语言(即由用 tff.federated_computation 装饰的 Python 函数的序列化主体构造的逻辑)被设计用作与平台无关的胶水语言。目前,此胶水语言用来从 TensorFlow 代码(限于 tff.tf_computation 块)的嵌入部分构建分布式系统。在时间充裕的情况下,我们会预见需要嵌入其他非 TensorFlow 逻辑的各个部分(例如可能表示输入流水线的关系数据库查询),它们全部使用相同的胶水语言(tff.federated_computation 块)相互连接。

TensorFlow 逻辑

声明 TensorFlow 计算

TFF 旨在配合 TensorFlow 使用。这样,您将在 TFF 中编写的大部分代码很可能是普通的(即本地执行的) TensorFlow 代码。为了在 TensorFlow 中使用此类代码,如上所述,只需用 tff.tf_computation 对其进行装饰。

例如,下面实现了一个函数,该函数接受数字并向其加 0.5

@tff.tf_computation(tf.float32)
def add_half(x):
  return tf.add(x, 0.5)

再次看到此内容,您可能想知道我们为什么应该定义另一个装饰器 tff.tf_computation,而不是简单地使用现有机制(如 tf.function)。与前一部分不同,我们在这里处理的是一个普通的 TensorFlow 代码块。

这样做有几个原因,虽然对它们的完整处理超出了本教程的范围,但下面这个主要原因值得注意:

  • 若要将使用 TensorFlow 代码实现的可重用构建块嵌入到联合计算的主体中,它们需要满足某些属性(例如在定义时进行跟踪和序列化、具有类型签名等)。这通常需要某种形式的装饰器。

一般而言,我们建议尽可能使用 TensorFlow 的原生机制进行组合(如 tf.function),因为 TFF 的装饰器与 Eager 函数进行交互的确切方式可能会逐步变化。

现在,回到上面的示例代码段,我们刚才定义的 add_half 计算可以像任何其他 TFF 计算一样由 TFF 处理。尤其是,它具有 TFF 类型签名。

str(add_half.type_signature)
'(float32 -> float32)'

请注意,此类型签名没有布局。TensorFlow 计算无法使用或返回联合类型。

现在,您还可以在其他计算中将 add_half 用作构建块。例如,下面是使用 tff.federated_map 算子在客户端设备上将 add_half 逐点应用到联合浮点的所有成员组成的方法。

@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def add_half_on_clients(x):
  return tff.federated_map(add_half, x)
str(add_half_on_clients.type_signature)
'({float32}@CLIENTS -> {float32}@CLIENTS)'

执行 TensorFlow 计算

执行使用 tff.tf_computation 定义的计算所遵循的规则与我们为 tff.federated_computation 描述的规则相同。可以将它们作为 Python 中的普通可调用对象进行调用,如下所示。

add_half_on_clients([1.0, 3.0, 2.0])
[<tf.Tensor: shape=(), dtype=float32, numpy=1.5>,
 <tf.Tensor: shape=(), dtype=float32, numpy=3.5>,
 <tf.Tensor: shape=(), dtype=float32, numpy=2.5>]

同样值得注意的是,以这种方式调用 add_half_on_clients 计算会模拟分布式过程。数据会在客户端上使用,并在客户端上返回。实际上,此计算会让每个客户端执行一次本地操作。此系统中没有明确提及 tff.SERVER(但在实践中,编排此类处理可能会用到)。可以将以这种方式定义的计算理解为在概念上类似于 MapReduce 中的 Map 阶段。

另外请记住,我们在前一个部分中讲的关于 TFF 计算会在定义时序列化的内容对 tff.tf_computation 代码同样适用,add_half_on_clients 的 Python 主体会在定义时被跟踪一次。在后续调用中,TFF 将使用其序列化后的表示形式。

tff.federated_computation 装饰的 Python 方法与用 tff.tf_computation 装饰的方法之间的唯一区别是,后者会被序列化为 TensorFlow 计算图(而不允许前者包含直接嵌入其中的 TensorFlow 代码)。

在后台,每个用 tff.tf_computation 装饰的方法会暂时停用 Eager Execution,以便捕获计算的结构。虽然 Eager Execution 已在本地停用,但只要您编写的计算逻辑能够正确序列化,欢迎您使用 Eager TensorFlow、AutoGraph、TensorFlow 2.0 构造等。

例如,以下代码将会失败:

try:

  # Eager mode
  constant_10 = tf.constant(10.)

  @tff.tf_computation(tf.float32)
  def add_ten(x):
    return x + constant_10

except Exception as err:
  print (err)
Attempting to capture an EagerTensor without building a function.

上述代码失败的原因是,constant_10 在计算图外部构造,而该计算图是 tff.tf_computation 在序列化过程中在 add_ten 的主体内部构造的。

另一方面,您可以对在 tff.tf_computation 内部调用时修改当前计算图的 Python 函数进行调用:

def get_constant_10():
  return tf.constant(10.)

@tff.tf_computation(tf.float32)
def add_ten(x):
  return x + get_constant_10()

add_ten(5.0)
15.0

请注意,TensorFlow 中的序列化机制正在逐步完善,我们期望 TFF 对计算进行序列化方式的细节也将逐步完善。

使用 tf.data.Dataset

如前所述,tff.tf_computation 的独特之处在于,它们允许您使用由您的代码作为形式参数抽象定义的 tf.data.Dataset。如果参数需要在 TensorFlow 中表示为数据集,则需要使用 tff.SequenceType 构造函数对其进行声明。

例如,类型规范 tff.SequenceType(tf.float32) 定义了 TFF 中浮点元素的抽象序列。序列可以包含张量或复杂的嵌套结构(稍后我们将看到相关示例)。T 型项的序列的简明表示形式为 T*

float32_sequence = tff.SequenceType(tf.float32)

str(float32_sequence)
'float32*'

假设在温度传感器示例中,每个传感器包含不只一个温度读数,而是多个。您可以使用下面的代码在 TensorFlow 中定义 TFF 计算,该计算将使用 tf.data.Dataset.reduce 算子在单个本地数据集中计算温度的平均值。

@tff.tf_computation(tff.SequenceType(tf.float32))
def get_local_temperature_average(local_temperatures):
  sum_and_count = (
      local_temperatures.reduce((0.0, 0), lambda x, y: (x[0] + y, x[1] + 1)))
  return sum_and_count[0] / tf.cast(sum_and_count[1], tf.float32)
str(get_local_temperature_average.type_signature)
'(float32* -> float32)'

在用 tff.tf_computation 装饰的方法的主体中,TFF 序列类型的形式参数简单地表示为行为类似 tf.data.Dataset 的对象(即支持相同的属性和方法,它们目前未作为该类型的子类实现,随着 TensorFlow 中对数据集的支持不断发展,这可能会发生变化)。

您可以轻松验证这一点,如下所示。

@tff.tf_computation(tff.SequenceType(tf.int32))
def foo(x):
  return x.reduce(np.int32(0), lambda x, y: x + y)

foo([1, 2, 3])
6

请记住,与普通的 tf.data.Dataset 不同,这些类似数据集的对象是占位符。它们不包含任何元素,因为它们表示抽象的序列类型参数,在具体上下文中使用时将绑定到具体数据。目前,对抽象定义的占位符数据的支持仍有一定局限,在早期的 TFF 中,您可能会遇到某些限制,但在本教程中不必担心这个问题(有关详细信息,请参阅文档页面)。

当在模拟模式下本地执行接受序列的计算时(如本教程所示),您可以将该序列作为 Python 列表进行馈送,如下所示(还可以用其他方式,例如,在 Eager 模式中作为 tf.data.Dataset 进行馈送,但现在我们将简单化处理)。

get_local_temperature_average([68.5, 70.3, 69.8])
69.53333

与其他 TFF 类型一样,上面定义的序列可以使用 tff.StructType 构造函数定义嵌套结构。例如,下面是一个声明计算的方法,该计算接受 AB 的序列对,并返回其乘积的和。我们将跟踪语句包含在计算主体中,以便您能够看到 TFF 类型签名如何转换为数据集的 output_typesoutput_shapes

@tff.tf_computation(tff.SequenceType(collections.OrderedDict([('A', tf.int32), ('B', tf.int32)])))
def foo(ds):
  print('element_structure = {}'.format(ds.element_spec))
  return ds.reduce(np.int32(0), lambda total, x: total + x['A'] * x['B'])
element_structure = OrderedDict([('A', TensorSpec(shape=(), dtype=tf.int32, name=None)), ('B', TensorSpec(shape=(), dtype=tf.int32, name=None))])

str(foo.type_signature)
'(<A=int32,B=int32>* -> int32)'
foo([{'A': 2, 'B': 3}, {'A': 4, 'B': 5}])
26

尽管将 tf.data.Datasets 用作形式参数在简单场景(如本教程中的场景)中有效,但对它的支持仍有局限且在不断发展。

汇总

现在,让我们再次尝试在联合设置中使用 TensorFlow 计算。假设我们有一组传感器,每个传感器有一个本地温度读数的序列。我们可以通过平均传感器的本地平均值来计算全局平均温度,如下所示。

@tff.federated_computation(
    tff.FederatedType(tff.SequenceType(tf.float32), tff.CLIENTS))
def get_global_temperature_average(sensor_readings):
  return tff.federated_mean(
      tff.federated_map(get_local_temperature_average, sensor_readings))

请注意,这并非是来自所有客户端的所有本地温度读数的简单平均,因为这需要根据不同客户端本地维护的读数数量权衡其贡献。我们将其作为练习留给读者来更新上面的代码;tff.federated_mean 算子接受权重作为可选的第二个参数(预计为联合浮点)。

还要注意,get_global_temperature_average 的输入现在变成了联合浮点序列。联合序列是我们通常在联合学习中表示设备端数据的方式,序列元素通常表示数据批次(稍后您将看到相关示例)。

str(get_global_temperature_average.type_signature)
'({float32*}@CLIENTS -> float32@SERVER)'

下面是如何用 Python 在数据样本上本地执行计算的方法。请注意,我们现在提供输入的方式是作为 listlist。外部列表在通过 tff.CLIENTS 表示的组中对设备进行迭代,内部列表在每个设备的本地序列中对元素进行迭代。

get_global_temperature_average([[68.0, 70.0], [71.0], [68.0, 72.0, 70.0]])
70.0

本教程的第一部分到此结束。我们鼓励您继续学习第二部分