提交 7c0121a5 编辑于 作者: Jason Roselander's avatar Jason Roselander 提交者: tensorflow-copybara
浏览文件

Creating TFF release 0.13.0

PiperOrigin-RevId: 300591979
上级 a8ea2945
...@@ -91,6 +91,7 @@ versa. ...@@ -91,6 +91,7 @@ versa.
TensorFlow Federated | TensorFlow TensorFlow Federated | TensorFlow
--------------------------------------------------------------------- | ---------- --------------------------------------------------------------------- | ----------
[0.13.0](https://github.com/tensorflow/federated/tree/v0.13.0) | [tensorflow 2.1.0](https://pypi.org/project/tensorflow/2.1.0/)
[0.12.0](https://github.com/tensorflow/federated/tree/v0.12.0) | [tensorflow 2.1.0](https://pypi.org/project/tensorflow/2.1.0/) [0.12.0](https://github.com/tensorflow/federated/tree/v0.12.0) | [tensorflow 2.1.0](https://pypi.org/project/tensorflow/2.1.0/)
[0.11.0](https://github.com/tensorflow/federated/tree/v0.11.0) | [tensorflow 2.0.0](https://pypi.org/project/tensorflow/2.0.0/) [0.11.0](https://github.com/tensorflow/federated/tree/v0.11.0) | [tensorflow 2.0.0](https://pypi.org/project/tensorflow/2.0.0/)
[0.10.1](https://github.com/tensorflow/federated/tree/v0.10.1) | [tensorflow 2.0.0](https://pypi.org/project/tensorflow/2.0.0/) [0.10.1](https://github.com/tensorflow/federated/tree/v0.10.1) | [tensorflow 2.0.0](https://pypi.org/project/tensorflow/2.0.0/)
......
# Release 0.13.0
## Major Features and Improvements
* Updated `absl-py` package dependency to `0.9.0`.
* Updated `h5py` package dependency to `2.8.0`.
* Updated `numpy` package dependency to `1.17.5`.
* Updated `tensorflow-privacy` package dependency to `0.2.2`.
## Breaking Changes
* Deprecated `dummy_batch` parameter of the `tff.learning.from_keras_model`
function.
## Bug Fixes
* Fixed issues with executor service using old executor API.
* Fixed issues with remote executor test using old executor API.
* Fixed issues in tutorial notebooks.
# Release 0.12.0 # Release 0.12.0
## Major Features and Improvements ## Major Features and Improvements
......
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
##### Copyright 2019 The TensorFlow Authors. ##### Copyright 2019 The TensorFlow Authors.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
#@title Licensed under the Apache License, Version 2.0 (the "License"); #@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
# You may obtain a copy of the License at # You may obtain a copy of the License at
# #
# https://www.apache.org/licenses/LICENSE-2.0 # https://www.apache.org/licenses/LICENSE-2.0
# #
# Unless required by applicable law or agreed to in writing, software # Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, # distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
# Custom Federated Algorithms, Part 1: Introduction to the Federated Core # Custom Federated Algorithms, Part 1: Introduction to the Federated Core
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
<table class="tfo-notebook-buttons" align="left"> <table class="tfo-notebook-buttons" align="left">
<td> <td>
<a target="_blank" href="https://www.tensorflow.org/federated/tutorials/custom_federated_algorithms_1"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a> <a target="_blank" href="https://www.tensorflow.org/federated/tutorials/custom_federated_algorithms_1"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
</td> </td>
<td> <td>
<a target="_blank" href="https://colab.research.google.com/github/tensorflow/federated/blob/v0.12.0/docs/tutorials/custom_federated_algorithms_1.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a> <a target="_blank" href="https://colab.research.google.com/github/tensorflow/federated/blob/v0.13.0/docs/tutorials/custom_federated_algorithms_1.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
</td> </td>
<td> <td>
<a target="_blank" href="https://github.com/tensorflow/federated/blob/v0.12.0/docs/tutorials/custom_federated_algorithms_1.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a> <a target="_blank" href="https://github.com/tensorflow/federated/blob/v0.13.0/docs/tutorials/custom_federated_algorithms_1.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
</td> </td>
</table> </table>
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
This tutorial is the first part of a two-part series that demonstrates how to This tutorial is the first part of a two-part series that demonstrates how to
implement custom types of federated algorithms in TensorFlow Federated (TFF) implement custom types of federated algorithms in TensorFlow Federated (TFF)
using the [Federated Core (FC)](../federated_core.md) - a set of lower-level using the [Federated Core (FC)](../federated_core.md) - a set of lower-level
interfaces that serve as a foundation upon which we have implemented the interfaces that serve as a foundation upon which we have implemented the
[Federated Learning (FL)](../federated_learning.md) layer. [Federated Learning (FL)](../federated_learning.md) layer.
This first part is more conceptual; we introduce some of the key concepts and This first part is more conceptual; we introduce some of the key concepts and
programming abstractions used in TFF, and we demonstrate their use on a very programming abstractions used in TFF, and we demonstrate their use on a very
simple example with a distributed array of temperature sensors. In simple example with a distributed array of temperature sensors. In
[the second part of this series](custom_federated_alrgorithms_2.ipynb), we use [the second part of this series](custom_federated_alrgorithms_2.ipynb), we use
the mechanisms we introduce here to implement a simple version of federated the mechanisms we introduce here to implement a simple version of federated
training and evaluation algorithms. As a follow-up, we encourage you to study training and evaluation algorithms. As a follow-up, we encourage you to study
[the implementation](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/federated_averaging.py) [the implementation](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/federated_averaging.py)
of federated averaging in `tff.learning`. of federated averaging in `tff.learning`.
By the end of this series, you should be able to recognize that the applications By the end of this series, you should be able to recognize that the applications
of Federated Core are not necessarily limited to learning. The programming of Federated Core are not necessarily limited to learning. The programming
abstractions we offer are quite generic, and could be used, e.g., to implement abstractions we offer are quite generic, and could be used, e.g., to implement
analytics and other custom types of computations over distributed data. analytics and other custom types of computations over distributed data.
Although this tutorial is designed to be self-contained, we encourage you to Although this tutorial is designed to be self-contained, we encourage you to
first read tutorials on first read tutorials on
[image classification](federated_learning_for_image_classification.ipynb) and [image classification](federated_learning_for_image_classification.ipynb) and
[text generation](federated_learning_for_text_generation.ipynb) for a [text generation](federated_learning_for_text_generation.ipynb) for a
higher-level and more gentle introduction to the TensorFlow Federated framework higher-level and more gentle introduction to the TensorFlow Federated framework
and the [Federated Learning](../federated_learning.md) APIs (`tff.learning`), as and the [Federated Learning](../federated_learning.md) APIs (`tff.learning`), as
it will help you put the concepts we describe here in context. it will help you put the concepts we describe here in context.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Intended Uses ## Intended Uses
In a nutshell, Federated Core (FC) is a development environment that makes it In a nutshell, Federated Core (FC) is a development environment that makes it
possible to compactly express program logic that combines TensorFlow code with possible to compactly express program logic that combines TensorFlow code with
distributed communication operators, such as those that are used in distributed communication operators, such as those that are used in
[Federated Averaging](https://arxiv.org/abs/1602.05629) - computing [Federated Averaging](https://arxiv.org/abs/1602.05629) - computing
distributed sums, averages, and other types of distributed aggregations over a distributed sums, averages, and other types of distributed aggregations over a
set of client devices in the system, broadcasting models and parameters to those set of client devices in the system, broadcasting models and parameters to those
devices, etc. devices, etc.
You may be aware of You may be aware of
[`tf.contrib.distribute`](https://www.tensorflow.org/api_docs/python/tf/contrib/distribute), [`tf.contrib.distribute`](https://www.tensorflow.org/api_docs/python/tf/contrib/distribute),
and a natural question to ask at this point may be: in what ways does this and a natural question to ask at this point may be: in what ways does this
framework differ? Both frameworks attempt at making TensorFlow computations framework differ? Both frameworks attempt at making TensorFlow computations
distributed, after all. distributed, after all.
One way to think about it is that, whereas the stated goal of One way to think about it is that, whereas the stated goal of
`tf.contrib.distribute` is *to allow users to use existing models and training `tf.contrib.distribute` is *to allow users to use existing models and training
code with minimal changes to enable distributed training*, and much focus is on code with minimal changes to enable distributed training*, and much focus is on
how to take advantage of distributed infrastructure to make existing training how to take advantage of distributed infrastructure to make existing training
code more efficient, the goal of TFF's Federated Core is to give researchers and code more efficient, the goal of TFF's Federated Core is to give researchers and
practitioners explicit control over the specific patterns of distributed practitioners explicit control over the specific patterns of distributed
communication they will use in their systems. The focus in FC is on providing a communication they will use in their systems. The focus in FC is on providing a
flexible and extensible language for expressing distributed data flow flexible and extensible language for expressing distributed data flow
algorithms, rather than a concrete set of implemented distributed training algorithms, rather than a concrete set of implemented distributed training
capabilities. capabilities.
One of the primary target audiences for TFF's FC API is researchers and One of the primary target audiences for TFF's FC API is researchers and
practitioners who might want to experiment with new federated learning practitioners who might want to experiment with new federated learning
algorithms and evaluate the consequences of subtle design choices that affect algorithms and evaluate the consequences of subtle design choices that affect
the manner in which the flow of data in the distributed system is orchestrated, the manner in which the flow of data in the distributed system is orchestrated,
yet without getting bogged down by system implementation details. The level of yet without getting bogged down by system implementation details. The level of
abstraction that FC API is aiming for roughly corresponds to pseudocode one abstraction that FC API is aiming for roughly corresponds to pseudocode one
could use to describe the mechanics of a federated learning algorithm in a could use to describe the mechanics of a federated learning algorithm in a
research publication - what data exists in the system and how it is transformed, research publication - what data exists in the system and how it is transformed,
but without dropping to the level of individual point-to-point network message but without dropping to the level of individual point-to-point network message
exchanges. exchanges.
TFF as a whole is targeting scenarios in which data is distributed, and must TFF as a whole is targeting scenarios in which data is distributed, and must
remain such, e.g., for privacy reasons, and where collecting all data at a remain such, e.g., for privacy reasons, and where collecting all data at a
centralized location may not be a viable option. This has implication on the centralized location may not be a viable option. This has implication on the
implementation of machine learning algorithms that require an increased degree implementation of machine learning algorithms that require an increased degree
of explicit control, as compared to scenarios in which all data can be of explicit control, as compared to scenarios in which all data can be
accumulated in a centralized location at a data center. accumulated in a centralized location at a data center.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Before we start ## Before we start
Before we dive into the code, please try to run the following "Hello World" Before we dive into the code, please try to run the following "Hello World"
example to make sure your environment is correctly setup. If it doesn't work, example to make sure your environment is correctly setup. If it doesn't work,
please refer to the [Installation](../install.md) guide for instructions. please refer to the [Installation](../install.md) guide for instructions.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
#@test {"skip": true} #@test {"skip": true}
!pip install --quiet --upgrade tensorflow_federated !pip install --quiet --upgrade tensorflow_federated
# Note: Jupyter requires a patch to asyncio. # Note: Jupyter requires a patch to asyncio.
!pip install --quiet --upgrade nest_asyncio !pip install --quiet --upgrade nest_asyncio
import nest_asyncio import nest_asyncio
nest_asyncio.apply() nest_asyncio.apply()
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
import collections import collections
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
import tensorflow_federated as tff import tensorflow_federated as tff
tf.compat.v1.enable_v2_behavior() tf.compat.v1.enable_v2_behavior()
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.federated_computation @tff.federated_computation
def hello_world(): def hello_world():
return 'Hello, World!' return 'Hello, World!'
hello_world() hello_world()
``` ```
%% Output %% Output
b'Hello, World!' b'Hello, World!'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Federated data ## Federated data
One of the distinguishing features of TFF is that it allows you to compactly One of the distinguishing features of TFF is that it allows you to compactly
express TensorFlow-based computations on *federated data*. We will be using the express TensorFlow-based computations on *federated data*. We will be using the
term *federated data* in this tutorial to refer to a collection of data items term *federated data* in this tutorial to refer to a collection of data items
hosted across a group of devices in a distributed system. For example, hosted across a group of devices in a distributed system. For example,
applications running on mobile devices may collect data and store it locally, applications running on mobile devices may collect data and store it locally,
without uploading to a centralized location. Or, an array of distributed sensors without uploading to a centralized location. Or, an array of distributed sensors
may collect and store temperature readings at their locations. may collect and store temperature readings at their locations.
Federated data like those in the above examples are treated in TFF as Federated data like those in the above examples are treated in TFF as
[first-class citizens](https://en.wikipedia.org/wiki/First-class_citizen), i.e., [first-class citizens](https://en.wikipedia.org/wiki/First-class_citizen), i.e.,
they may appear as parameters and results of functions, and they have types. To they may appear as parameters and results of functions, and they have types. To
reinforce this notion, we will refer to federated data sets as *federated reinforce this notion, we will refer to federated data sets as *federated
values*, or as *values of federated types*. values*, or as *values of federated types*.
The important point to understand is that we are modeling the entire collection The important point to understand is that we are modeling the entire collection
of data items across all devices (e.g., the entire collection temperature of data items across all devices (e.g., the entire collection temperature
readings from all sensors in a distributed array) as a single federated value. readings from all sensors in a distributed array) as a single federated value.
For example, here's how one would define in TFF the type of a *federated float* For example, here's how one would define in TFF the type of a *federated float*
hosted by a group of client devices. A collection of temperature readings that hosted by a group of client devices. A collection of temperature readings that
materialize across an array of distributed sensors could be modeled as a value materialize across an array of distributed sensors could be modeled as a value
of this federated type. of this federated type.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
federated_float_on_clients = tff.FederatedType(tf.float32, tff.CLIENTS) federated_float_on_clients = tff.FederatedType(tf.float32, tff.CLIENTS)
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
More generally, a federated type in TFF is defined by specifying the type `T` of More generally, a federated type in TFF is defined by specifying the type `T` of
its *member constituents* - the items of data that reside on individual devices, its *member constituents* - the items of data that reside on individual devices,
and the group `G` of devices on which federated values of this type are hosted and the group `G` of devices on which federated values of this type are hosted
(plus a third, optional bit of information we'll mention shortly). We refer to (plus a third, optional bit of information we'll mention shortly). We refer to
the group `G` of devices hosting a federated value as the value's *placement*. the group `G` of devices hosting a federated value as the value's *placement*.
Thus, `tff.CLIENTS` is an example of a placement. Thus, `tff.CLIENTS` is an example of a placement.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(federated_float_on_clients.member) str(federated_float_on_clients.member)
``` ```
%% Output %% Output
'float32' 'float32'
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(federated_float_on_clients.placement) str(federated_float_on_clients.placement)
``` ```
%% Output %% Output
'CLIENTS' 'CLIENTS'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
A federated type with member constituents `T` and placement `G` can be A federated type with member constituents `T` and placement `G` can be
represented compactly as `{T}@G`, as shown below. represented compactly as `{T}@G`, as shown below.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(federated_float_on_clients) str(federated_float_on_clients)
``` ```
%% Output %% Output
'{float32}@CLIENTS' '{float32}@CLIENTS'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
The curly braces `{}` in this concise notation serve as a reminder that the The curly braces `{}` in this concise notation serve as a reminder that the
member constituents (items of data on different devices) may differ, as you member constituents (items of data on different devices) may differ, as you
would expect e.g., of temperature sensor readings, so the clients as a group are would expect e.g., of temperature sensor readings, so the clients as a group are
jointly hosting a [multi-set](https://en.wikipedia.org/wiki/Multiset) of jointly hosting a [multi-set](https://en.wikipedia.org/wiki/Multiset) of
`T`-typed items that together constitute the federated value. `T`-typed items that together constitute the federated value.
It is important to note that the member constituents of a federated value are It is important to note that the member constituents of a federated value are
generally opaque to the programmer, i.e., a federated value should not be generally opaque to the programmer, i.e., a federated value should not be
thought of as a simple `dict` keyed by an identifier of a device in the system - thought of as a simple `dict` keyed by an identifier of a device in the system -
these values are intended to be collectively transformed only by *federated these values are intended to be collectively transformed only by *federated
operators* that abstractly represent various kinds of distributed communication operators* that abstractly represent various kinds of distributed communication
protocols (such as aggregation). If this sounds too abstract, don't worry - we protocols (such as aggregation). If this sounds too abstract, don't worry - we
will return to this shortly, and we will illustrate it with concrete examples. will return to this shortly, and we will illustrate it with concrete examples.
Federated types in TFF come in two flavors: those where the member constituents Federated types in TFF come in two flavors: those where the member constituents
of a federated value may differ (as just seen above), and those where they are of a federated value may differ (as just seen above), and those where they are
known to be all equal. This is controlled by the third, optional `all_equal` known to be all equal. This is controlled by the third, optional `all_equal`
parameter in the `tff.FederatedType` constructor (defaulting to `False`). parameter in the `tff.FederatedType` constructor (defaulting to `False`).
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
federated_float_on_clients.all_equal federated_float_on_clients.all_equal
``` ```
%% Output %% Output
False False
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
A federated type with a placement `G` in which all of the `T`-typed member A federated type with a placement `G` in which all of the `T`-typed member
constituents are known to be equal can be compactly represented as `T@G` (as constituents are known to be equal can be compactly represented as `T@G` (as
opposed to `{T}@G`, that is, with the curly braces dropped to reflect the fact opposed to `{T}@G`, that is, with the curly braces dropped to reflect the fact
that the multi-set of member constituents consists of a single item). that the multi-set of member constituents consists of a single item).
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(tff.FederatedType(tf.float32, tff.CLIENTS, all_equal=True)) str(tff.FederatedType(tf.float32, tff.CLIENTS, all_equal=True))
``` ```
%% Output %% Output
'float32@CLIENTS' 'float32@CLIENTS'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
One example of a federated value of such type that might arise in practical One example of a federated value of such type that might arise in practical
scenarios is a hyperparameter (such as a learning rate, a clipping norm, etc.) scenarios is a hyperparameter (such as a learning rate, a clipping norm, etc.)
that has been broadcasted by a server to a group of devices that participate in that has been broadcasted by a server to a group of devices that participate in
federated training. federated training.
Another example is a set of parameters for a machine learning model pre-trained Another example is a set of parameters for a machine learning model pre-trained
at the server, that were then broadcasted to a group of client devices, where at the server, that were then broadcasted to a group of client devices, where
they can be personalized for each user. they can be personalized for each user.
For example, suppose we have a pair of `float32` parameters `a` and `b` for a For example, suppose we have a pair of `float32` parameters `a` and `b` for a
simple one-dimensional linear regression model. We can construct the simple one-dimensional linear regression model. We can construct the
(non-federated) type of such models for use in TFF as follows. The angle braces (non-federated) type of such models for use in TFF as follows. The angle braces
`<>` in the printed type string are a compact TFF notation for named or unnamed `<>` in the printed type string are a compact TFF notation for named or unnamed
tuples. tuples.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
simple_regression_model_type = ( simple_regression_model_type = (
tff.NamedTupleType([('a', tf.float32), ('b', tf.float32)])) tff.NamedTupleType([('a', tf.float32), ('b', tf.float32)]))
str(simple_regression_model_type) str(simple_regression_model_type)
``` ```
%% Output %% Output
'<a=float32,b=float32>' '<a=float32,b=float32>'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Note that we are only specifying `dtype`s above. Non-scalar types are also Note that we are only specifying `dtype`s above. Non-scalar types are also
supported. In the above code, `tf.float32` is a shortcut notation for the more supported. In the above code, `tf.float32` is a shortcut notation for the more
general `tff.TensorType(dtype=tf.float32, shape=[])`. general `tff.TensorType(dtype=tf.float32, shape=[])`.
When this model is broadcasted to clients, the type of the resulting federated When this model is broadcasted to clients, the type of the resulting federated
value can be represented as shown below. value can be represented as shown below.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(tff.FederatedType( str(tff.FederatedType(
simple_regression_model_type, tff.CLIENTS, all_equal=True)) simple_regression_model_type, tff.CLIENTS, all_equal=True))
``` ```
%% Output %% Output
'<a=float32,b=float32>@CLIENTS' '<a=float32,b=float32>@CLIENTS'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Per symmetry with *federated float* above, we will refer to such a type as a Per symmetry with *federated float* above, we will refer to such a type as a
*federated tuple*. More generally, we'll often use the term *federated XYZ* to *federated tuple*. More generally, we'll often use the term *federated XYZ* to
refer to a federated value in which member constituents are *XYZ*-like. Thus, we refer to a federated value in which member constituents are *XYZ*-like. Thus, we
will talk about things like *federated tuples*, *federated sequences*, will talk about things like *federated tuples*, *federated sequences*,
*federated models*, and so on. *federated models*, and so on.
Now, coming back to `float32@CLIENTS` - while it appears replicated across Now, coming back to `float32@CLIENTS` - while it appears replicated across
multiple devices, it is actually a single `float32`, since all member are the multiple devices, it is actually a single `float32`, since all member are the
same. In general, you may think of any *all-equal* federated type, i.e., one of same. In general, you may think of any *all-equal* federated type, i.e., one of
the form `T@G`, as isomorphic to a non-federated type `T`, since in both cases, the form `T@G`, as isomorphic to a non-federated type `T`, since in both cases,
there's actually only a single (albeit potentially replicated) item of type `T`. there's actually only a single (albeit potentially replicated) item of type `T`.
Given the isomorphism between `T` and `T@G`, you may wonder what purpose, if Given the isomorphism between `T` and `T@G`, you may wonder what purpose, if
any, the latter types might serve. Read on. any, the latter types might serve. Read on.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Placements ## Placements
### Design Overview ### Design Overview
In the preceding section, we've introduced the concept of *placements* - groups In the preceding section, we've introduced the concept of *placements* - groups
of system participants that might be jointly hosting a federated value, and of system participants that might be jointly hosting a federated value, and
we've demonstrated the use of `tff.CLIENTS` as an example specification of a we've demonstrated the use of `tff.CLIENTS` as an example specification of a
placement. placement.
To explain why the notion of a *placement* is so fundamental that we needed to To explain why the notion of a *placement* is so fundamental that we needed to
incorporate it into the TFF type system, recall what we mentioned at the incorporate it into the TFF type system, recall what we mentioned at the
beginning of this tutorial about some of the intended uses of TFF. beginning of this tutorial about some of the intended uses of TFF.
Although in this tutorial, you will only see TFF code being executed locally in Although in this tutorial, you will only see TFF code being executed locally in
a simulated environment, our goal is for TFF to enable writing code that you a simulated environment, our goal is for TFF to enable writing code that you
could deploy for execution on groups of physical devices in a distributed could deploy for execution on groups of physical devices in a distributed
system, potentially including mobile or embedded devices running Android. Each system, potentially including mobile or embedded devices running Android. Each
of of those devices would receive a separate set of instructions to execute of of those devices would receive a separate set of instructions to execute
locally, depending on the role it plays in the system (an end-user device, a locally, depending on the role it plays in the system (an end-user device, a
centralized coordinator, an intermediate layer in a multi-tier architecture, centralized coordinator, an intermediate layer in a multi-tier architecture,
etc.). It is important to be able to reason about which subsets of devices etc.). It is important to be able to reason about which subsets of devices
execute what code, and where different portions of the data might physically execute what code, and where different portions of the data might physically
materialize. materialize.
This is especially important when dealing with, e.g., application data on mobile This is especially important when dealing with, e.g., application data on mobile
devices. Since the data is private and can be sensitive, we need the ability to devices. Since the data is private and can be sensitive, we need the ability to
statically verify that this data will never leave the device (and prove facts statically verify that this data will never leave the device (and prove facts
about how the data is being processed). The placement specifications are one of about how the data is being processed). The placement specifications are one of
the mechanisms designed to support this. the mechanisms designed to support this.
TFF has been designed as a data-centric programming environment, and as such, TFF has been designed as a data-centric programming environment, and as such,
unlike some of the existing frameworks that focus on *operations* and where unlike some of the existing frameworks that focus on *operations* and where
those operations might *run*, TFF focuses on *data*, where that data those operations might *run*, TFF focuses on *data*, where that data
*materializes*, and how it's being *transformed*. Consequently, placement is *materializes*, and how it's being *transformed*. Consequently, placement is
modeled as a property of data in TFF, rather than as a property of operations on modeled as a property of data in TFF, rather than as a property of operations on
data. Indeed, as you're about to see in the next section, some of the TFF data. Indeed, as you're about to see in the next section, some of the TFF
operations span across locations, and run "in the network", so to speak, rather operations span across locations, and run "in the network", so to speak, rather
than being executed by a single machine or a group of machines. than being executed by a single machine or a group of machines.
Representing the type of a certain value as `T@G` or `{T}@G` (as opposed to just Representing the type of a certain value as `T@G` or `{T}@G` (as opposed to just
`T`) makes data placement decisions explicit, and together with a static `T`) makes data placement decisions explicit, and together with a static
analysis of programs written in TFF, it can serve as a foundation for providing analysis of programs written in TFF, it can serve as a foundation for providing
formal privacy guarantees for sensitive on-device data. formal privacy guarantees for sensitive on-device data.
An important thing to note at this point, however, is that while we encourage An important thing to note at this point, however, is that while we encourage
TFF users to be explicit about *groups* of participating devices that host the TFF users to be explicit about *groups* of participating devices that host the
data (the placements), the programmer will never deal with the raw data or data (the placements), the programmer will never deal with the raw data or
identities of the *individual* participants. identities of the *individual* participants.
(Note: While it goes far outside the scope of this tutorial, we should mention (Note: While it goes far outside the scope of this tutorial, we should mention
that there is one notable exception to the above, a `tff.federated_collect` that there is one notable exception to the above, a `tff.federated_collect`
operator that is intended as a low-level primitive, only for specialized operator that is intended as a low-level primitive, only for specialized
situations. Its explicit use in situations where it can be avoided is not situations. Its explicit use in situations where it can be avoided is not
recommended, as it may limit the possible future applications. For example, if recommended, as it may limit the possible future applications. For example, if
during the course of static analysis, we determine that a computation uses such during the course of static analysis, we determine that a computation uses such
low-level mechanisms, we may disallow its access to certain types of data.) low-level mechanisms, we may disallow its access to certain types of data.)
Within the body of TFF code, by design, there's no way to enumerate the devices Within the body of TFF code, by design, there's no way to enumerate the devices
that constitute the group represented by `tff.CLIENTS`, or to probe for the that constitute the group represented by `tff.CLIENTS`, or to probe for the
existence of a specific device in the group. There's no concept of a device or existence of a specific device in the group. There's no concept of a device or
client identity anywhere in the Federated Core API, the underlying set of client identity anywhere in the Federated Core API, the underlying set of
architectural abstractions, or the core runtime infrastructure we provide to architectural abstractions, or the core runtime infrastructure we provide to
support simulations. All the computation logic you write will be expressed as support simulations. All the computation logic you write will be expressed as
operations on the entire client group. operations on the entire client group.
Recall here what we mentioned earlier about values of federated types being Recall here what we mentioned earlier about values of federated types being
unlike Python `dict`, in that one cannot simply enumerate their member unlike Python `dict`, in that one cannot simply enumerate their member
constituents. Think of values that your TFF program logic manipulates as being constituents. Think of values that your TFF program logic manipulates as being
associated with placements (groups), rather than with individual participants. associated with placements (groups), rather than with individual participants.
Placements *are* designed to be a first-class citizen in TFF as well, and can Placements *are* designed to be a first-class citizen in TFF as well, and can
appear as parameters and results of a `placement` type (to be represented by appear as parameters and results of a `placement` type (to be represented by
`tff.PlacementType` in the API). In the future, we plan to provide a variety of `tff.PlacementType` in the API). In the future, we plan to provide a variety of
operators to transform or combine placements, but this is outside the scope of operators to transform or combine placements, but this is outside the scope of
this tutorial. For now, it suffices to think of `placement` as an opaque this tutorial. For now, it suffices to think of `placement` as an opaque
primitive built-in type in TFF, similar to how `int` and `bool` are opaque primitive built-in type in TFF, similar to how `int` and `bool` are opaque
built-in types in Python, with `tff.CLIENTS` being a constant literal of this built-in types in Python, with `tff.CLIENTS` being a constant literal of this
type, not unlike `1` being a constant literal of type `int`. type, not unlike `1` being a constant literal of type `int`.
### Specifying Placements ### Specifying Placements
TFF provides two basic placement literals, `tff.CLIENTS` and `tff.SERVER`, to TFF provides two basic placement literals, `tff.CLIENTS` and `tff.SERVER`, to
make it easy to express the rich variety of practical scenarios that are make it easy to express the rich variety of practical scenarios that are
naturally modeled as client-server architectures, with multiple *client* devices naturally modeled as client-server architectures, with multiple *client* devices
(mobile phones, embedded devices, distributed databases, sensors, etc.) (mobile phones, embedded devices, distributed databases, sensors, etc.)
orchestrated by a single centralized *server* coordinator. TFF is designed to orchestrated by a single centralized *server* coordinator. TFF is designed to
also support custom placements, multiple client groups, multi-tiered and other, also support custom placements, multiple client groups, multi-tiered and other,
more general distributed architectures, but discussing them is outside the scope more general distributed architectures, but discussing them is outside the scope
of this tutorial. of this tutorial.
TFF doesn't prescribe what either the `tff.CLIENTS` or the `tff.SERVER` actually TFF doesn't prescribe what either the `tff.CLIENTS` or the `tff.SERVER` actually
represent. represent.
In particular, `tff.SERVER` may be a single physical device (a member of a In particular, `tff.SERVER` may be a single physical device (a member of a
singleton group), but it might just as well be a group of replicas in a singleton group), but it might just as well be a group of replicas in a
fault-tolerant cluster running state machine replication - we do not make any fault-tolerant cluster running state machine replication - we do not make any
special architectural assumptions. Rather, we use the `all_equal` bit mentioned special architectural assumptions. Rather, we use the `all_equal` bit mentioned
in the preceding section to express the fact that we're generally dealing with in the preceding section to express the fact that we're generally dealing with
only a single item of data at the server. only a single item of data at the server.
Likewise, `tff.CLIENTS` in some applications might represent all clients in the Likewise, `tff.CLIENTS` in some applications might represent all clients in the
system - what in the context of federated learning we sometimes refer to as the system - what in the context of federated learning we sometimes refer to as the
*population*, but e.g., in *population*, but e.g., in
[production implementations of Federated Averaging](https://arxiv.org/abs/1602.05629), [production implementations of Federated Averaging](https://arxiv.org/abs/1602.05629),
it may represent a *cohort* - a subset of the clients selected for paticipation it may represent a *cohort* - a subset of the clients selected for paticipation
in a particular round of training. The abstractly defined placements are given in a particular round of training. The abstractly defined placements are given
concrete meaning when a computation in which they appear is deployed for concrete meaning when a computation in which they appear is deployed for
execution (or simply invoked like a Python function in a simulated environment, execution (or simply invoked like a Python function in a simulated environment,
as is demonstrated in this tutorial). In our local simulations, the group of as is demonstrated in this tutorial). In our local simulations, the group of
clients is determined by the federated data supplied as input. clients is determined by the federated data supplied as input.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Federated computations ## Federated computations
### Declaring federated computations ### Declaring federated computations
TFF is designed as a strongly-typed functional programming environment that TFF is designed as a strongly-typed functional programming environment that
supports modular development. supports modular development.
The basic unit of composition in TFF is a *federated computation* - a section of The basic unit of composition in TFF is a *federated computation* - a section of
logic that may accept federated values as input and return federated values as logic that may accept federated values as input and return federated values as
output. Here's how you can define a computation that calculates the average of output. Here's how you can define a computation that calculates the average of
the temperatures reported by the sensor array from our previous example. the temperatures reported by the sensor array from our previous example.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS)) @tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(sensor_readings): def get_average_temperature(sensor_readings):
return tff.federated_mean(sensor_readings) return tff.federated_mean(sensor_readings)
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Looking at the above code, at this point you might be asking - aren't there Looking at the above code, at this point you might be asking - aren't there
already decorator constructs to define composable units such as already decorator constructs to define composable units such as
[`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function) [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function)
in TensorFlow, and if so, why introduce yet another one, and how is it in TensorFlow, and if so, why introduce yet another one, and how is it
different? different?
The short answer is that the code generated by the `tff.federated_computation` The short answer is that the code generated by the `tff.federated_computation`
wrapper is *neither* TensorFlow, *nor is it* Python - it's a specification of a wrapper is *neither* TensorFlow, *nor is it* Python - it's a specification of a
distributed system in an internal platform-independent *glue* language. At this distributed system in an internal platform-independent *glue* language. At this
point, this will undoubtedly sound cryptic, but please bear this intuitive point, this will undoubtedly sound cryptic, but please bear this intuitive
interpretation of a federated computation as an abstract specification of a interpretation of a federated computation as an abstract specification of a
distributed system in mind. We'll explain it in a minute. distributed system in mind. We'll explain it in a minute.
First, let's play with the definition a bit. TFF computations are generally First, let's play with the definition a bit. TFF computations are generally
modeled as functions - with or without parameters, but with well-defined type modeled as functions - with or without parameters, but with well-defined type
signatures. You can print the type signature of a computation by querying its signatures. You can print the type signature of a computation by querying its
`type_signature` property, as shown below. `type_signature` property, as shown below.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(get_average_temperature.type_signature) str(get_average_temperature.type_signature)
``` ```
%% Output %% Output
'({float32}@CLIENTS -> float32@SERVER)' '({float32}@CLIENTS -> float32@SERVER)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
The type signature tells us that the computation accepts a collection of The type signature tells us that the computation accepts a collection of
different sensor readings on client devices, and returns a single average on the different sensor readings on client devices, and returns a single average on the
server. server.
Before we go any further, let's reflect on this for a minute - the input and Before we go any further, let's reflect on this for a minute - the input and
output of this computation are *in different places* (on `CLIENTS` vs. at the output of this computation are *in different places* (on `CLIENTS` vs. at the
`SERVER`). Recall what we said in the preceding section on placements about how `SERVER`). Recall what we said in the preceding section on placements about how
*TFF operations may span across locations, and run in the network*, and what we *TFF operations may span across locations, and run in the network*, and what we
just said about federated computations as representing abstract specifications just said about federated computations as representing abstract specifications
of distributed systems. We have just a defined one such computation - a simple of distributed systems. We have just a defined one such computation - a simple
distributed system in which data is consumed at client devices, and the distributed system in which data is consumed at client devices, and the
aggregate results emerge at the server. aggregate results emerge at the server.
In many practical scenarios, the computations that represent top-level tasks In many practical scenarios, the computations that represent top-level tasks
will tend to accept their inputs and report their outputs at the server - this will tend to accept their inputs and report their outputs at the server - this
reflects the idea that computations might be triggered by *queries* that reflects the idea that computations might be triggered by *queries* that
originate and terminate on the server. originate and terminate on the server.
However, FC API does not impose this assumption, and many of the building blocks However, FC API does not impose this assumption, and many of the building blocks
we use internally (including numerous `tff.federated_...` operators you may find we use internally (including numerous `tff.federated_...` operators you may find
in the API) have inputs and outputs with distinct placements, so in general, you in the API) have inputs and outputs with distinct placements, so in general, you
should not think about a federated computation as something that *runs on the should not think about a federated computation as something that *runs on the
server* or is *executed by a server*. The server is just one type of participant server* or is *executed by a server*. The server is just one type of participant
in a federated computation. In thinking about the mechanics of such in a federated computation. In thinking about the mechanics of such
computations, it's best to always default to the global network-wide computations, it's best to always default to the global network-wide
perspective, rather than the perspective of a single centralized coordinator. perspective, rather than the perspective of a single centralized coordinator.
In general, functional type signatures are compactly represented as `(T -> U)` In general, functional type signatures are compactly represented as `(T -> U)`
for types `T` and `U` of inputs and outputs, respectively. The type of the for types `T` and `U` of inputs and outputs, respectively. The type of the
formal parameter (such `sensor_readings` in this case) is specified as the formal parameter (such `sensor_readings` in this case) is specified as the
argument to the decorator. You don't need to specify the type of the result - argument to the decorator. You don't need to specify the type of the result -
it's determined automatically. it's determined automatically.
Although TFF does offer limited forms of polymorphism, programmers are strongly Although TFF does offer limited forms of polymorphism, programmers are strongly
encouraged to be explicit about the types of data they work with, as that makes encouraged to be explicit about the types of data they work with, as that makes
understanding, debugging, and formally verifying properties of your code easier. understanding, debugging, and formally verifying properties of your code easier.
In some cases, explicitly specifying types is a requirement (e.g., polymorphic In some cases, explicitly specifying types is a requirement (e.g., polymorphic
computations are currently not directly executable). computations are currently not directly executable).
### Executing federated computations ### Executing federated computations
In order to support development and debugging, TFF allows you to directly invoke In order to support development and debugging, TFF allows you to directly invoke
computations defined this way as Python functions, as shown below. Where the computations defined this way as Python functions, as shown below. Where the
computation expects a value of a federated type with the `all_equal` bit set to computation expects a value of a federated type with the `all_equal` bit set to
`False`, you can feed it as a plain `list` in Python, and for federated types `False`, you can feed it as a plain `list` in Python, and for federated types
with the `all_equal` bit set to `True`, you can just directly feed the (single) with the `all_equal` bit set to `True`, you can just directly feed the (single)
member constituent. This is also how the results are reported back to you. member constituent. This is also how the results are reported back to you.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
get_average_temperature([68.5, 70.3, 69.8]) get_average_temperature([68.5, 70.3, 69.8])
``` ```
%% Output %% Output
69.53334 69.53334
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
When running computations like this in simulation mode, you act as an external When running computations like this in simulation mode, you act as an external
observer with a system-wide view, who has the ability to supply inputs and observer with a system-wide view, who has the ability to supply inputs and
consume outputs at any locations in the network, as indeed is the case here - consume outputs at any locations in the network, as indeed is the case here -
you supplied client values at input, and consumed the server result. you supplied client values at input, and consumed the server result.
Now, let's return to a note we made earlier about the Now, let's return to a note we made earlier about the
`tff.federated_computation` decorator emitting code in a *glue* language. `tff.federated_computation` decorator emitting code in a *glue* language.
Although the logic of TFF computations can be expressed as ordinary functions in Although the logic of TFF computations can be expressed as ordinary functions in
Python (you just need to decorate them with `tff.federated_computation` as we've Python (you just need to decorate them with `tff.federated_computation` as we've
done above), and you can directly invoke them with Python arguments just done above), and you can directly invoke them with Python arguments just
like any other Python functions in this notebook, behind the scenes, as we noted like any other Python functions in this notebook, behind the scenes, as we noted
earlier, TFF computations are actually *not* Python. earlier, TFF computations are actually *not* Python.
What we mean by this is that when the Python interpreter encounters a function What we mean by this is that when the Python interpreter encounters a function
decorated with `tff.federated_computation`, it traces the statements in this decorated with `tff.federated_computation`, it traces the statements in this
function's body once (at definition time), and then constructs a function's body once (at definition time), and then constructs a
[serialized representation](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/proto/v0/computation.proto) [serialized representation](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/proto/v0/computation.proto)
of the computation's logic for future use - whether for execution, or to be of the computation's logic for future use - whether for execution, or to be
incorporated as a sub-component into another computation. incorporated as a sub-component into another computation.
You can verify this by adding a print statement, as follows: You can verify this by adding a print statement, as follows:
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS)) @tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(sensor_readings): def get_average_temperature(sensor_readings):
print ('Getting traced, the argument is "{}".'.format( print ('Getting traced, the argument is "{}".'.format(
type(sensor_readings).__name__)) type(sensor_readings).__name__))
return tff.federated_mean(sensor_readings) return tff.federated_mean(sensor_readings)
``` ```
%% Output %% Output
Getting traced, the argument is "ValueImpl". Getting traced, the argument is "ValueImpl".
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
You can think of Python code that defines a federated computation similarly to You can think of Python code that defines a federated computation similarly to
how you would think of Python code that builds a TensorFlow graph in a non-eager how you would think of Python code that builds a TensorFlow graph in a non-eager
context (if you're not familiar with the non-eager uses of TensorFlow, think of context (if you're not familiar with the non-eager uses of TensorFlow, think of
your Python code defining a graph of operations to be executed later, but not your Python code defining a graph of operations to be executed later, but not
actually running them on the fly). The non-eager graph-building code in actually running them on the fly). The non-eager graph-building code in
TensorFlow is Python, but the TensorFlow graph constructed by this code is TensorFlow is Python, but the TensorFlow graph constructed by this code is
platform-independent and serializable. platform-independent and serializable.
Likewise, TFF computations are defined in Python, but the Python statements in Likewise, TFF computations are defined in Python, but the Python statements in
their bodies, such as `tff.federated_mean` in the example weve just shown, their bodies, such as `tff.federated_mean` in the example weve just shown,
are compiled into a portable and platform-independent serializable are compiled into a portable and platform-independent serializable
representation under the hood. representation under the hood.
As a developer, you don't need to concern yourself with the details of this As a developer, you don't need to concern yourself with the details of this
representation, as you will never need to directly work with it, but you should representation, as you will never need to directly work with it, but you should
be aware of its existence, the fact that TFF computations are fundamentally be aware of its existence, the fact that TFF computations are fundamentally
non-eager, and cannot capture arbitrary Python state. Python code contained in a non-eager, and cannot capture arbitrary Python state. Python code contained in a
TFF computation's body is executed at definition time, when the body of the TFF computation's body is executed at definition time, when the body of the
Python function decorated with `tff.federated_computation` is traced before Python function decorated with `tff.federated_computation` is traced before
getting serialized. It's not retraced again at invocation time (except when the getting serialized. It's not retraced again at invocation time (except when the
function is polymorphic; please refer to the documentation pages for details). function is polymorphic; please refer to the documentation pages for details).
You may wonder why we've chosen to introduce a dedicated internal non-Python You may wonder why we've chosen to introduce a dedicated internal non-Python
representation. One reason is that ultimately, TFF computations are intended to representation. One reason is that ultimately, TFF computations are intended to
be deployable to real physical environments, and hosted on mobile or embedded be deployable to real physical environments, and hosted on mobile or embedded
devices, where Python may not be available. devices, where Python may not be available.
Another reason is that TFF computations express the global behavior of Another reason is that TFF computations express the global behavior of
distributed systems, as opposed to Python programs which express the local distributed systems, as opposed to Python programs which express the local
behavior of individual participants. You can see that in the simple example behavior of individual participants. You can see that in the simple example
above, with the special operator `tff.federated_mean` that accepts data on above, with the special operator `tff.federated_mean` that accepts data on
client devices, but deposits the results on the server. client devices, but deposits the results on the server.
The operator `tff.federated_mean` cannot be easily modeled as an ordinary The operator `tff.federated_mean` cannot be easily modeled as an ordinary
operator in Python, since it doesn't execute locally - as noted earlier, it operator in Python, since it doesn't execute locally - as noted earlier, it
represents a distributed system that coordinates the behavior of multiple system represents a distributed system that coordinates the behavior of multiple system
participants. We will refer to such operators as *federated operators*, to participants. We will refer to such operators as *federated operators*, to
distinguish them from ordinary (local) operators in Python. distinguish them from ordinary (local) operators in Python.
The TFF type system, and the fundamental set of operations supported in the TFF's The TFF type system, and the fundamental set of operations supported in the TFF's
language, thus deviates significantly from those in Python, necessitating the language, thus deviates significantly from those in Python, necessitating the
use of a dedicated representation. use of a dedicated representation.
### Composing federated computations ### Composing federated computations
As noted above, federated computations and their constituents are best As noted above, federated computations and their constituents are best
understood as models of distributed systems, and you can think of composing understood as models of distributed systems, and you can think of composing
federated computations as composing more complex distributed systems from federated computations as composing more complex distributed systems from
simpler ones. You can think of the `tff.federated_mean` operator as a kind of simpler ones. You can think of the `tff.federated_mean` operator as a kind of
built-in template federated computation with a type signature `({T}@CLIENTS -> built-in template federated computation with a type signature `({T}@CLIENTS ->
T@SERVER)` (indeed, just like computations you write, this operator also has a T@SERVER)` (indeed, just like computations you write, this operator also has a
complex structure - under the hood we break it down into simpler operators). complex structure - under the hood we break it down into simpler operators).
The same is true of composing federated computations. The computation The same is true of composing federated computations. The computation
`get_average_temperature` may be invoked in a body of another Python function `get_average_temperature` may be invoked in a body of another Python function
decorated with `tff.federated_computation` - doing so will cause it to be decorated with `tff.federated_computation` - doing so will cause it to be
embedded in the body of the parent, much in the same way `tff.federated_mean` embedded in the body of the parent, much in the same way `tff.federated_mean`
was embedded in its own body earlier. was embedded in its own body earlier.
An important restriction to be aware of is that bodies of Python functions An important restriction to be aware of is that bodies of Python functions
decorated with `tff.federated_computation` must consist *only* of federated decorated with `tff.federated_computation` must consist *only* of federated
operators, i.e., they cannot directly contain TensorFlow operations. For operators, i.e., they cannot directly contain TensorFlow operations. For
example, you cannot directly use `tf.nest` interfaces to add a pair of example, you cannot directly use `tf.nest` interfaces to add a pair of
federated values. TensorFlow code must be confined to blocks of code decorated federated values. TensorFlow code must be confined to blocks of code decorated
with a `tff.tf_computation` discussed in the following section. Only when with a `tff.tf_computation` discussed in the following section. Only when
wrapped in this manner can the wrapped TensorFlow code be invoked in the body of wrapped in this manner can the wrapped TensorFlow code be invoked in the body of
a `tff.federated_computation`. a `tff.federated_computation`.
The reasons for this separation are technical (it's hard to trick operators such The reasons for this separation are technical (it's hard to trick operators such
as `tf.add` to work with non-tensors) as well as architectural. The language of as `tf.add` to work with non-tensors) as well as architectural. The language of
federated computations (i.e., the logic constructed from serialized bodies of federated computations (i.e., the logic constructed from serialized bodies of
Python functions decorated with `tff.federated_computation`) is designed to Python functions decorated with `tff.federated_computation`) is designed to
serve as a platform-independent *glue* language. This glue language is currently serve as a platform-independent *glue* language. This glue language is currently
used to build distributed systems from embedded sections of TensorFlow code used to build distributed systems from embedded sections of TensorFlow code
(confined to `tff.tf_computation` blocks). In the fullness of time, we (confined to `tff.tf_computation` blocks). In the fullness of time, we
anticipate the need to embed sections of other, non-TensorFlow logic, such as anticipate the need to embed sections of other, non-TensorFlow logic, such as
relational database queries that might represent input pipelines, all connected relational database queries that might represent input pipelines, all connected
together using the same glue language (the `tff.federated_computation` blocks). together using the same glue language (the `tff.federated_computation` blocks).
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## TensorFlow logic ## TensorFlow logic
### Declaring TensorFlow computations ### Declaring TensorFlow computations
TFF is designed for use with TensorFlow. As such, the bulk of the code you will TFF is designed for use with TensorFlow. As such, the bulk of the code you will
write in TFF is likely to be ordinary (i.e., locally-executing) TensorFlow code. write in TFF is likely to be ordinary (i.e., locally-executing) TensorFlow code.
In order to use such code with TFF, as noted above, it just needs to be In order to use such code with TFF, as noted above, it just needs to be
decorated with `tff.tf_computation`. decorated with `tff.tf_computation`.
For example, here's how we could implement a function that takes a number and For example, here's how we could implement a function that takes a number and
adds `0.5` to it. adds `0.5` to it.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.tf_computation(tf.float32) @tff.tf_computation(tf.float32)
def add_half(x): def add_half(x):
return tf.add(x, 0.5) return tf.add(x, 0.5)
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Once again, looking at this, you may be wondering why we should define another Once again, looking at this, you may be wondering why we should define another
decorator `tff.tf_computation` instead of simply using an existing mechanism decorator `tff.tf_computation` instead of simply using an existing mechanism
such as `tf.function`. Unlike in the preceding section, here we are such as `tf.function`. Unlike in the preceding section, here we are
dealing with an ordinary block of TensorFlow code. dealing with an ordinary block of TensorFlow code.
There are a few reasons for this, the full treatment of which goes beyond the There are a few reasons for this, the full treatment of which goes beyond the
scope of this tutorial, but it's worth naming the main one: scope of this tutorial, but it's worth naming the main one:
* In order to embed reusable building blocks implemented using TensorFlow code * In order to embed reusable building blocks implemented using TensorFlow code
in the bodies of federated computations, they need to satisfy certain in the bodies of federated computations, they need to satisfy certain
properties - such as getting traced and serialized at definition time, properties - such as getting traced and serialized at definition time,
having type signatures, etc. This generally requires some form of a having type signatures, etc. This generally requires some form of a
decorator. decorator.
In general, we recommend using TensorFlow's native mechanisms for composition, In general, we recommend using TensorFlow's native mechanisms for composition,
such as `tf.function`, wherever possible, as the exact manner in such as `tf.function`, wherever possible, as the exact manner in
which TFF's decorator interacts with eager functions can be expected to evolve. which TFF's decorator interacts with eager functions can be expected to evolve.
Now, coming back to the example code snippet above, the computation `add_half` Now, coming back to the example code snippet above, the computation `add_half`
we just defined can be treated by TFF just like any other TFF computation. In we just defined can be treated by TFF just like any other TFF computation. In
particular, it has a TFF type signature. particular, it has a TFF type signature.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(add_half.type_signature) str(add_half.type_signature)
``` ```
%% Output %% Output
'(float32 -> float32)' '(float32 -> float32)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Note this type signature does not have placements. TensorFlow computations Note this type signature does not have placements. TensorFlow computations
cannot consume or return federated types. cannot consume or return federated types.
You can now also use `add_half` as a building block in other computations . For You can now also use `add_half` as a building block in other computations . For
example, here's how you can use the `tff.federated_map` operator to apply example, here's how you can use the `tff.federated_map` operator to apply
`add_half` pointwise to all member constituents of a federated float on client `add_half` pointwise to all member constituents of a federated float on client
devices. devices.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS)) @tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def add_half_on_clients(x): def add_half_on_clients(x):
return tff.federated_map(add_half, x) return tff.federated_map(add_half, x)
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(add_half_on_clients.type_signature) str(add_half_on_clients.type_signature)
``` ```
%% Output %% Output
'({float32}@CLIENTS -> {float32}@CLIENTS)' '({float32}@CLIENTS -> {float32}@CLIENTS)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Executing TensorFlow computations ### Executing TensorFlow computations
Execution of computations defined with `tff.tf_computation` follows the same Execution of computations defined with `tff.tf_computation` follows the same
rules as those we described for `tff.federated_computation`. They can be invoked rules as those we described for `tff.federated_computation`. They can be invoked
as ordinary callables in Python, as follows. as ordinary callables in Python, as follows.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
add_half_on_clients([1.0, 3.0, 2.0]) add_half_on_clients([1.0, 3.0, 2.0])
``` ```
%% Output %% Output
[<tf.Tensor: shape=(), dtype=float32, numpy=1.5>, [<tf.Tensor: shape=(), dtype=float32, numpy=1.5>,
<tf.Tensor: shape=(), dtype=float32, numpy=3.5>, <tf.Tensor: shape=(), dtype=float32, numpy=3.5>,
<tf.Tensor: shape=(), dtype=float32, numpy=2.5>] <tf.Tensor: shape=(), dtype=float32, numpy=2.5>]
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Once again, it is worth noting that invoking the computation Once again, it is worth noting that invoking the computation
`add_half_on_clients` in this manner simulates a distirbuted process. Data is `add_half_on_clients` in this manner simulates a distirbuted process. Data is
consumed on clients, and returned on clients. Indeed, this computation has each consumed on clients, and returned on clients. Indeed, this computation has each
client perform a local action. There is no `tff.SERVER` explicitly mentioned in client perform a local action. There is no `tff.SERVER` explicitly mentioned in
this system (even if in practice, orchestrating such processing might involve this system (even if in practice, orchestrating such processing might involve
one). Think of a computation defined this way as conceptually analogous to the one). Think of a computation defined this way as conceptually analogous to the
`Map` stage in `MapReduce`. `Map` stage in `MapReduce`.
Also, keep in mind that what we said in the preceding section about TFF Also, keep in mind that what we said in the preceding section about TFF
computations getting serialized at the definition time remains true for computations getting serialized at the definition time remains true for
`tff.tf_computation` code as well - the Python body of `add_half_on_clients` `tff.tf_computation` code as well - the Python body of `add_half_on_clients`
gets traced once at definition time. On subsequent invocations, TFF uses its gets traced once at definition time. On subsequent invocations, TFF uses its
serialized representation. serialized representation.
The only difference between Python methods decorated with The only difference between Python methods decorated with
`tff.federated_computation` and those decorated with `tff.tf_computation` is `tff.federated_computation` and those decorated with `tff.tf_computation` is
that the latter are serialized as TensorFlow graphs (whereas the former are not that the latter are serialized as TensorFlow graphs (whereas the former are not
allowed to contain TensorFlow code directly embedded in them). allowed to contain TensorFlow code directly embedded in them).
Under the hood, each method decorated with `tff.tf_computation` temporarily Under the hood, each method decorated with `tff.tf_computation` temporarily
disables eager execution in order to allow the computation's structure to be disables eager execution in order to allow the computation's structure to be
captured. While eager execution is locally disabled, you are welcome to use captured. While eager execution is locally disabled, you are welcome to use
eager TensorFlow, AutoGraph, TensorFlow 2.0 constructs, etc., so long as you eager TensorFlow, AutoGraph, TensorFlow 2.0 constructs, etc., so long as you
write the logic of your computation in a manner such that it can get correctly write the logic of your computation in a manner such that it can get correctly
serialized. serialized.
For example, the following code will fail: For example, the following code will fail:
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
try: try:
# Eager mode # Eager mode
constant_10 = tf.constant(10.) constant_10 = tf.constant(10.)
@tff.tf_computation(tf.float32) @tff.tf_computation(tf.float32)
def add_ten(x): def add_ten(x):
return x + constant_10 return x + constant_10
except Exception as err: except Exception as err:
print (err) print (err)
``` ```
%% Output %% Output
Attempting to capture an EagerTensor without building a function. Attempting to capture an EagerTensor without building a function.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
The above fails because `constant_10` has already been constructed outside of The above fails because `constant_10` has already been constructed outside of
the graph that `tff.tf_computation` constructs internally in the body of the graph that `tff.tf_computation` constructs internally in the body of
`add_ten` during the serialization process. `add_ten` during the serialization process.
On the other hand, invoking python functions that modify the current graph when On the other hand, invoking python functions that modify the current graph when
called inside a `tff.tf_computation` is fine: called inside a `tff.tf_computation` is fine:
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
def get_constant_10(): def get_constant_10():
return tf.constant(10.) return tf.constant(10.)
@tff.tf_computation(tf.float32) @tff.tf_computation(tf.float32)
def add_ten(x): def add_ten(x):
return x + get_constant_10() return x + get_constant_10()
add_ten(5.0) add_ten(5.0)
``` ```
%% Output %% Output
15.0 15.0
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Note that the serialization mechanisms in TensorFlow are evolving, and we expect Note that the serialization mechanisms in TensorFlow are evolving, and we expect
the details of how TFF serializes computations to evolve as well. the details of how TFF serializes computations to evolve as well.
### Working with `tf.data.Dataset`s ### Working with `tf.data.Dataset`s
As noted earlier, a unique feature of `tff.tf_computation`s is that they allows As noted earlier, a unique feature of `tff.tf_computation`s is that they allows
you to work with `tf.data.Dataset`s defined abstractly as formal parameters by you to work with `tf.data.Dataset`s defined abstractly as formal parameters by
your code. Parameters to be represented in TensorFlow as data sets need to be your code. Parameters to be represented in TensorFlow as data sets need to be
declared using the `tff.SequenceType` constructor. declared using the `tff.SequenceType` constructor.
For example, the type specification `tff.SequenceType(tf.float32)` defines an For example, the type specification `tff.SequenceType(tf.float32)` defines an
abstract sequence of float elements in TFF. Sequences can contain either abstract sequence of float elements in TFF. Sequences can contain either
tensors, or complex nested structures (we'll see examples of those later). The tensors, or complex nested structures (we'll see examples of those later). The
concise representation of a sequence of `T`-typed items is `T*`. concise representation of a sequence of `T`-typed items is `T*`.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
float32_sequence = tff.SequenceType(tf.float32) float32_sequence = tff.SequenceType(tf.float32)
str(float32_sequence) str(float32_sequence)
``` ```
%% Output %% Output
'float32*' 'float32*'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Suppose that in our temperature sensor example, each sensor holds not just one Suppose that in our temperature sensor example, each sensor holds not just one
temperature reading, but multiple. Here's how you can define a TFF computation temperature reading, but multiple. Here's how you can define a TFF computation
in TensorFlow that calculates the average of temperatures in a single local data in TensorFlow that calculates the average of temperatures in a single local data
set using the `tf.data.Dataset.reduce` operator. set using the `tf.data.Dataset.reduce` operator.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.tf_computation(tff.SequenceType(tf.float32)) @tff.tf_computation(tff.SequenceType(tf.float32))
def get_local_temperature_average(local_temperatures): def get_local_temperature_average(local_temperatures):
sum_and_count = ( sum_and_count = (
local_temperatures.reduce((0.0, 0), lambda x, y: (x[0] + y, x[1] + 1))) 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) return sum_and_count[0] / tf.cast(sum_and_count[1], tf.float32)
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(get_local_temperature_average.type_signature) str(get_local_temperature_average.type_signature)
``` ```
%% Output %% Output
'(float32* -> float32)' '(float32* -> float32)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
In the body of a method decorated with `tff.tf_computation`, formal parameters In the body of a method decorated with `tff.tf_computation`, formal parameters
of a TFF sequence type are represented simply as objects that behave like of a TFF sequence type are represented simply as objects that behave like
`tf.data.Dataset`, i.e., support the same properties and methods (they are `tf.data.Dataset`, i.e., support the same properties and methods (they are
currently not implemented as subclasses of that type - this may change as the currently not implemented as subclasses of that type - this may change as the
support for data sets in TensorFlow evolves). support for data sets in TensorFlow evolves).
You can easily verify this as follows. You can easily verify this as follows.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.tf_computation(tff.SequenceType(tf.int32)) @tff.tf_computation(tff.SequenceType(tf.int32))
def foo(x): def foo(x):
return x.reduce(np.int32(0), lambda x, y: x + y) return x.reduce(np.int32(0), lambda x, y: x + y)
foo([1, 2, 3]) foo([1, 2, 3])
``` ```
%% Output %% Output
6 6
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Keep in mind that unlike ordinary `tf.data.Dataset`s, these dataset-like objects Keep in mind that unlike ordinary `tf.data.Dataset`s, these dataset-like objects
are placeholders. They don't contain any elements, since they represent abstract are placeholders. They don't contain any elements, since they represent abstract
sequence-typed parameters, to be bound to concrete data when used in a concrete sequence-typed parameters, to be bound to concrete data when used in a concrete
context. Support for abstractly-defined placeholder data sets is still somewhat context. Support for abstractly-defined placeholder data sets is still somewhat
limited at this point, and in the early days of TFF, you may encounter certain limited at this point, and in the early days of TFF, you may encounter certain
restrictions, but we won't need to worry about them in this tutorial (please restrictions, but we won't need to worry about them in this tutorial (please
refer to the documentation pages for details). refer to the documentation pages for details).
When locally executing a computation that accepts a sequence in a simulation When locally executing a computation that accepts a sequence in a simulation
mode, such as in this tutorial, you can feed the sequence as Python list, as mode, such as in this tutorial, you can feed the sequence as Python list, as
below (as well as in other ways, e.g., as a `tf.data.Dataset` in eager mode, but below (as well as in other ways, e.g., as a `tf.data.Dataset` in eager mode, but
for now, we'll keep it simple). for now, we'll keep it simple).
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
get_local_temperature_average([68.5, 70.3, 69.8]) get_local_temperature_average([68.5, 70.3, 69.8])
``` ```
%% Output %% Output
69.53333 69.53333
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Like all other TFF types, sequences like those defined above can use the Like all other TFF types, sequences like those defined above can use the
`tff.NamedTupleType` constructor to define nested structures. For example, `tff.NamedTupleType` constructor to define nested structures. For example,
here's how one could declare a computation that accepts a sequence of pairs `A`, here's how one could declare a computation that accepts a sequence of pairs `A`,
`B`, and returns the sum of their products. We include the tracing statements in `B`, and returns the sum of their products. We include the tracing statements in
the body of the computation so that you can see how the TFF type signature the body of the computation so that you can see how the TFF type signature
translates into the dataset's `output_types` and `output_shapes`. translates into the dataset's `output_types` and `output_shapes`.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.tf_computation(tff.SequenceType(collections.OrderedDict([('A', tf.int32), ('B', tf.int32)]))) @tff.tf_computation(tff.SequenceType(collections.OrderedDict([('A', tf.int32), ('B', tf.int32)])))
def foo(ds): def foo(ds):
print('element_structure = {}'.format(ds.element_spec)) print('element_structure = {}'.format(ds.element_spec))
return ds.reduce(np.int32(0), lambda total, x: total + x['A'] * x['B']) return ds.reduce(np.int32(0), lambda total, x: total + x['A'] * x['B'])
``` ```
%% Output %% Output
element_structure = OrderedDict([('A', TensorSpec(shape=(), dtype=tf.int32, name=None)), ('B', TensorSpec(shape=(), dtype=tf.int32, name=None))]) element_structure = OrderedDict([('A', TensorSpec(shape=(), dtype=tf.int32, name=None)), ('B', TensorSpec(shape=(), dtype=tf.int32, name=None))])
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(foo.type_signature) str(foo.type_signature)
``` ```
%% Output %% Output
'(<A=int32,B=int32>* -> int32)' '(<A=int32,B=int32>* -> int32)'
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
foo([{'A': 2, 'B': 3}, {'A': 4, 'B': 5}]) foo([{'A': 2, 'B': 3}, {'A': 4, 'B': 5}])
``` ```
%% Output %% Output
26 26
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
The support for using `tf.data.Datasets` as formal parameters is still somewhat The support for using `tf.data.Datasets` as formal parameters is still somewhat
limited and evolving, although functional in simple scenarios such as those used limited and evolving, although functional in simple scenarios such as those used
in this tutorial. in this tutorial.
## Putting it all together ## Putting it all together
Now, let's try again to use our TensorFlow computation in a federated setting. Now, let's try again to use our TensorFlow computation in a federated setting.
Suppose we have a group of sensors that each have a local sequence of Suppose we have a group of sensors that each have a local sequence of
temperature readings. We can compute the global temperature average by averaging temperature readings. We can compute the global temperature average by averaging
the sensors' local averages as follows. the sensors' local averages as follows.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.federated_computation( @tff.federated_computation(
tff.FederatedType(tff.SequenceType(tf.float32), tff.CLIENTS)) tff.FederatedType(tff.SequenceType(tf.float32), tff.CLIENTS))
def get_global_temperature_average(sensor_readings): def get_global_temperature_average(sensor_readings):
return tff.federated_mean( return tff.federated_mean(
tff.federated_map(get_local_temperature_average, sensor_readings)) tff.federated_map(get_local_temperature_average, sensor_readings))
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Note that this isn't a simple average across all local temperature readings from Note that this isn't a simple average across all local temperature readings from
all clients, as that would require weighing contributions from different clients all clients, as that would require weighing contributions from different clients
by the number of readings they locally maintain. We leave it as an exercise for by the number of readings they locally maintain. We leave it as an exercise for
the reader to update the above code; the `tff.federated_mean` operator the reader to update the above code; the `tff.federated_mean` operator
accepts the weight as an optional second argument (expected to be a federated accepts the weight as an optional second argument (expected to be a federated
float). float).
Also note that the input to `get_global_temperature_average` now becomes a Also note that the input to `get_global_temperature_average` now becomes a
*federated float sequence*. Federated sequences is how we will typically represent *federated float sequence*. Federated sequences is how we will typically represent
on-device data in federated learning, with sequence elements typically on-device data in federated learning, with sequence elements typically
representing data batches (you will see examples of this shortly). representing data batches (you will see examples of this shortly).
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(get_global_temperature_average.type_signature) str(get_global_temperature_average.type_signature)
``` ```
%% Output %% Output
'({float32*}@CLIENTS -> float32@SERVER)' '({float32*}@CLIENTS -> float32@SERVER)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Here's how we can locally execute the computation on a sample of data in Python. Here's how we can locally execute the computation on a sample of data in Python.
Notice that the way we supply the input is now as a `list` of `list`s. The outer Notice that the way we supply the input is now as a `list` of `list`s. The outer
list iterates over the devices in the group represented by `tff.CLIENTS`, and list iterates over the devices in the group represented by `tff.CLIENTS`, and
the inner ones iterate over elements in each device's local sequence. the inner ones iterate over elements in each device's local sequence.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
get_global_temperature_average([[68.0, 70.0], [71.0], [68.0, 72.0, 70.0]]) get_global_temperature_average([[68.0, 70.0], [71.0], [68.0, 72.0, 70.0]])
``` ```
%% Output %% Output
70.0 70.0
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
This concludes the first part of the tutorial... we encourage you to continue on This concludes the first part of the tutorial... we encourage you to continue on
to the [second part](custom_federated_algorithms_2.ipynb). to the [second part](custom_federated_algorithms_2.ipynb).
......
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
##### Copyright 2019 The TensorFlow Authors. ##### Copyright 2019 The TensorFlow Authors.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
#@title Licensed under the Apache License, Version 2.0 (the "License"); #@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
# You may obtain a copy of the License at # You may obtain a copy of the License at
# #
# https://www.apache.org/licenses/LICENSE-2.0 # https://www.apache.org/licenses/LICENSE-2.0
# #
# Unless required by applicable law or agreed to in writing, software # Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, # distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
# Custom Federated Algorithms, Part 2: Implementing Federated Averaging # Custom Federated Algorithms, Part 2: Implementing Federated Averaging
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
<table class="tfo-notebook-buttons" align="left"> <table class="tfo-notebook-buttons" align="left">
<td> <td>
<a target="_blank" href="https://www.tensorflow.org/federated/tutorials/custom_federated_algorithms_2"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a> <a target="_blank" href="https://www.tensorflow.org/federated/tutorials/custom_federated_algorithms_2"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
</td> </td>
<td> <td>
<a target="_blank" href="https://colab.research.google.com/github/tensorflow/federated/blob/v0.12.0/docs/tutorials/custom_federated_algorithms_2.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a> <a target="_blank" href="https://colab.research.google.com/github/tensorflow/federated/blob/v0.13.0/docs/tutorials/custom_federated_algorithms_2.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
</td> </td>
<td> <td>
<a target="_blank" href="https://github.com/tensorflow/federated/blob/v0.12.0/docs/tutorials/custom_federated_algorithms_2.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a> <a target="_blank" href="https://github.com/tensorflow/federated/blob/v0.13.0/docs/tutorials/custom_federated_algorithms_2.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
</td> </td>
</table> </table>
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
This tutorial is the second part of a two-part series that demonstrates how to This tutorial is the second part of a two-part series that demonstrates how to
implement custom types of federated algorithms in TFF using the implement custom types of federated algorithms in TFF using the
[Federated Core (FC)](../federated_core.md), which serves as a foundation for [Federated Core (FC)](../federated_core.md), which serves as a foundation for
the [Federated Learning (FL)](../federated_learning.md) layer (`tff.learning`). the [Federated Learning (FL)](../federated_learning.md) layer (`tff.learning`).
We encourage you to first read the We encourage you to first read the
[first part of this series](custom_federated_algorithms_1.ipynb), which [first part of this series](custom_federated_algorithms_1.ipynb), which
introduce some of the key concepts and programming abstractions used here. introduce some of the key concepts and programming abstractions used here.
This second part of the series uses the mechanisms introduced in the first part This second part of the series uses the mechanisms introduced in the first part
to implement a simple version of federated training and evaluation algorithms. to implement a simple version of federated training and evaluation algorithms.
We encourage you to review the We encourage you to review the
[image classification](federated_learning_for_image_classification.ipynb) and [image classification](federated_learning_for_image_classification.ipynb) and
[text generation](federated_learning_for_text_generation.ipynb) tutorials for a [text generation](federated_learning_for_text_generation.ipynb) tutorials for a
higher-level and more gentle introduction to TFF's Federated Learning APIs, as higher-level and more gentle introduction to TFF's Federated Learning APIs, as
they will help you put the concepts we describe here in context. they will help you put the concepts we describe here in context.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Before we start ## Before we start
Before we start, try to run the following "Hello World" example to make sure Before we start, try to run the following "Hello World" example to make sure
your environment is correctly setup. If it doesn't work, please refer to the your environment is correctly setup. If it doesn't work, please refer to the
[Installation](../install.md) guide for instructions. [Installation](../install.md) guide for instructions.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
#@test {"skip": true} #@test {"skip": true}
!pip install --quiet --upgrade tensorflow_federated !pip install --quiet --upgrade tensorflow_federated
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
import collections import collections
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
import tensorflow_federated as tff import tensorflow_federated as tff
tf.compat.v1.enable_v2_behavior() tf.compat.v1.enable_v2_behavior()
# TODO(b/148678573,b/148685415): must use the ReferenceExecutor because it # TODO(b/148678573,b/148685415): must use the ReferenceExecutor because it
# supports unbounded references and tff.sequence_* intrinsics. # supports unbounded references and tff.sequence_* intrinsics.
tff.framework.set_default_executor(tff.framework.ReferenceExecutor()) tff.framework.set_default_executor(tff.framework.ReferenceExecutor())
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.federated_computation @tff.federated_computation
def hello_world(): def hello_world():
return 'Hello, World!' return 'Hello, World!'
hello_world() hello_world()
``` ```
%% Output %% Output
'Hello, World!' 'Hello, World!'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Implementing Federated Averaging ## Implementing Federated Averaging
As in As in
[Federated Learning for Image Classification](federated_learning_for_image_classification.ipynb), [Federated Learning for Image Classification](federated_learning_for_image_classification.ipynb),
we are going to use the MNIST example, but since this is intended as a low-level we are going to use the MNIST example, but since this is intended as a low-level
tutorial, we are going to bypass the Keras API and `tff.simulation`, write raw tutorial, we are going to bypass the Keras API and `tff.simulation`, write raw
model code, and construct a federated data set from scratch. model code, and construct a federated data set from scratch.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Preparing federated data sets ### Preparing federated data sets
For the sake of a demonstration, we're going to simulate a scenario in which we For the sake of a demonstration, we're going to simulate a scenario in which we
have data from 10 users, and each of the users contributes knowledge how to have data from 10 users, and each of the users contributes knowledge how to
recognize a different digit. This is about as recognize a different digit. This is about as
non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables) non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables)
as it gets. as it gets.
First, let's load the standard MNIST data: First, let's load the standard MNIST data:
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
mnist_train, mnist_test = tf.keras.datasets.mnist.load_data() mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
[(x.dtype, x.shape) for x in mnist_train] [(x.dtype, x.shape) for x in mnist_train]
``` ```
%% Output %% Output
[(dtype('uint8'), (60000, 28, 28)), (dtype('uint8'), (60000,))] [(dtype('uint8'), (60000, 28, 28)), (dtype('uint8'), (60000,))]
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
The data comes as Numpy arrays, one with images and another with digit labels, both The data comes as Numpy arrays, one with images and another with digit labels, both
with the first dimension going over the individual examples. Let's write a with the first dimension going over the individual examples. Let's write a
helper function that formats it in a way compatible with how we feed federated helper function that formats it in a way compatible with how we feed federated
sequences into TFF computations, i.e., as a list of lists - the outer list sequences into TFF computations, i.e., as a list of lists - the outer list
ranging over the users (digits), the inner ones ranging over batches of data in ranging over the users (digits), the inner ones ranging over batches of data in
each client's sequence. As is customary, we will structure each batch as a pair each client's sequence. As is customary, we will structure each batch as a pair
of tensors named `x` and `y`, each with the leading batch dimension. While at of tensors named `x` and `y`, each with the leading batch dimension. While at
it, we'll also flatten each image into a 784-element vector and rescale the it, we'll also flatten each image into a 784-element vector and rescale the
pixels in it into the `0..1` range, so that we don't have to clutter the model pixels in it into the `0..1` range, so that we don't have to clutter the model
logic with data conversions. logic with data conversions.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
NUM_EXAMPLES_PER_USER = 1000 NUM_EXAMPLES_PER_USER = 1000
BATCH_SIZE = 100 BATCH_SIZE = 100
def get_data_for_digit(source, digit): def get_data_for_digit(source, digit):
output_sequence = [] output_sequence = []
all_samples = [i for i, d in enumerate(source[1]) if d == digit] 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): for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):
batch_samples = all_samples[i:i + BATCH_SIZE] batch_samples = all_samples[i:i + BATCH_SIZE]
output_sequence.append({ output_sequence.append({
'x': 'x':
np.array([source[0][i].flatten() / 255.0 for i in batch_samples], np.array([source[0][i].flatten() / 255.0 for i in batch_samples],
dtype=np.float32), dtype=np.float32),
'y': 'y':
np.array([source[1][i] for i in batch_samples], dtype=np.int32) np.array([source[1][i] for i in batch_samples], dtype=np.int32)
}) })
return output_sequence return output_sequence
federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)] 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)] federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
As a quick sanity check, let's look at the `Y` tensor in the last batch of data As a quick sanity check, let's look at the `Y` tensor in the last batch of data
contributed by the fifth client (the one corresponding to the digit `5`). contributed by the fifth client (the one corresponding to the digit `5`).
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
federated_train_data[5][-1]['y'] federated_train_data[5][-1]['y']
``` ```
%% Output %% Output
array([5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, array([5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5], dtype=int32) 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5], dtype=int32)
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Just to be sure, let's also look at the image corresponding to the last element of that batch. Just to be sure, let's also look at the image corresponding to the last element of that batch.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
from matplotlib import pyplot as plt from matplotlib import pyplot as plt
plt.imshow(federated_train_data[5][-1]['x'][-1].reshape(28, 28), cmap='gray') plt.imshow(federated_train_data[5][-1]['x'][-1].reshape(28, 28), cmap='gray')
plt.grid(False) plt.grid(False)
plt.show() plt.show()
``` ```
%% Output %% Output
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### On combining TensorFlow and TFF ### On combining TensorFlow and TFF
In this tutorial, for compactness we immediately decorate functions that In this tutorial, for compactness we immediately decorate functions that
introduce TensorFlow logic with `tff.tf_computation`. However, for more complex introduce TensorFlow logic with `tff.tf_computation`. However, for more complex
logic, this is not the pattern we recommend. Debugging TensorFlow can already be logic, this is not the pattern we recommend. Debugging TensorFlow can already be
a challenge, and debugging TensorFlow after it has been fully serialized and a challenge, and debugging TensorFlow after it has been fully serialized and
then re-imported necessarily loses some metadata and limits interactivity, then re-imported necessarily loses some metadata and limits interactivity,
making debugging even more of a challenge. making debugging even more of a challenge.
Therefore, **we strongly recommend writing complex TF logic as stand-alone Therefore, **we strongly recommend writing complex TF logic as stand-alone
Python functions** (that is, without `tff.tf_computation` decoration). This way Python functions** (that is, without `tff.tf_computation` decoration). This way
the TensorFlow logic can be developed and tested using TF best practices and the TensorFlow logic can be developed and tested using TF best practices and
tools (like eager mode), before serializing the computation for TFF (e.g., by invoking `tff.tf_computation` with a Python function as the argument). tools (like eager mode), before serializing the computation for TFF (e.g., by invoking `tff.tf_computation` with a Python function as the argument).
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Defining a loss function ### Defining a loss function
Now that we have the data, let's define a loss function that we can use for Now that we have the data, let's define a loss function that we can use for
training. First, let's define the type of input as a TFF named tuple. Since the training. First, let's define the type of input as a TFF named tuple. Since the
size of data batches may vary, we set the batch dimension to `None` to indicate size of data batches may vary, we set the batch dimension to `None` to indicate
that the size of this dimension is unknown. that the size of this dimension is unknown.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
BATCH_SPEC = collections.OrderedDict( BATCH_SPEC = collections.OrderedDict(
x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32), x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),
y=tf.TensorSpec(shape=[None], dtype=tf.int32)) y=tf.TensorSpec(shape=[None], dtype=tf.int32))
BATCH_TYPE = tff.to_type(BATCH_SPEC) BATCH_TYPE = tff.to_type(BATCH_SPEC)
str(BATCH_TYPE) str(BATCH_TYPE)
``` ```
%% Output %% Output
'<x=float32[?,784],y=int32[?]>' '<x=float32[?,784],y=int32[?]>'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
You may be wondering why we can't just define an ordinary Python type. Recall You may be wondering why we can't just define an ordinary Python type. Recall
the discussion in [part 1](custom_federated_algorithms_1.ipynb), where we the discussion in [part 1](custom_federated_algorithms_1.ipynb), where we
explained that while we can express the logic of TFF computations using Python, explained that while we can express the logic of TFF computations using Python,
under the hood TFF computations *are not* Python. The symbol `BATCH_TYPE` under the hood TFF computations *are not* Python. The symbol `BATCH_TYPE`
defined above represents an abstract TFF type specification. It is important to defined above represents an abstract TFF type specification. It is important to
distinguish this *abstract* TFF type from concrete Python *representation* distinguish this *abstract* TFF type from concrete Python *representation*
types, e.g., containers such as `dict` or `collections.namedtuple` that may be types, e.g., containers such as `dict` or `collections.namedtuple` that may be
used to represent the TFF type in the body of a Python function. Unlike Python, used to represent the TFF type in the body of a Python function. Unlike Python,
TFF has a single abstract type constructor `tff.NamedTupleType` for tuple-like TFF has a single abstract type constructor `tff.NamedTupleType` for tuple-like
containers, with elements that can be individually named or left unnamed. This containers, with elements that can be individually named or left unnamed. This
type is also used to model formal parameters of computations, as TFF type is also used to model formal parameters of computations, as TFF
computations can formally only declare one parameter and one result - you will computations can formally only declare one parameter and one result - you will
see examples of this shortly. see examples of this shortly.
Let's now define the TFF type of model parameters, again as a TFF named tuple of Let's now define the TFF type of model parameters, again as a TFF named tuple of
*weights* and *bias*. *weights* and *bias*.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
MODEL_SPEC = collections.OrderedDict( MODEL_SPEC = collections.OrderedDict(
weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32), weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),
bias=tf.TensorSpec(shape=[10], dtype=tf.float32)) bias=tf.TensorSpec(shape=[10], dtype=tf.float32))
MODEL_TYPE = tff.to_type(MODEL_SPEC) MODEL_TYPE = tff.to_type(MODEL_SPEC)
print(MODEL_TYPE) print(MODEL_TYPE)
``` ```
%% Output %% Output
<weights=float32[784,10],bias=float32[10]> <weights=float32[784,10],bias=float32[10]>
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
With those definitions in place, now we can define the loss for a given model, over a single batch. Note the usage of `@tf.function` decorator inside the `@tff.tf_computation` decorator. This allows us to write TF using Python like semantics even though were inside a `tf.Graph` context created by the `tff.tf_computation` decorator. With those definitions in place, now we can define the loss for a given model, over a single batch. Note the usage of `@tf.function` decorator inside the `@tff.tf_computation` decorator. This allows us to write TF using Python like semantics even though were inside a `tf.Graph` context created by the `tff.tf_computation` decorator.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE) @tff.tf_computation(MODEL_TYPE, BATCH_TYPE)
@tf.function @tf.function
def batch_loss(model, batch): def batch_loss(model, batch):
predicted_y = tf.nn.softmax( predicted_y = tf.nn.softmax(
tf.matmul(batch['x'], model['weights']) + model['bias']) tf.matmul(batch['x'], model['weights']) + model['bias'])
return -tf.reduce_mean( return -tf.reduce_mean(
tf.reduce_sum( tf.reduce_sum(
tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1])) tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
As expected, computation `batch_loss` returns `float32` loss given the model and As expected, computation `batch_loss` returns `float32` loss given the model and
a single data batch. Note how the `MODEL_TYPE` and `BATCH_TYPE` have been lumped a single data batch. Note how the `MODEL_TYPE` and `BATCH_TYPE` have been lumped
together into a 2-tuple of formal parameters; you can recognize the type of together into a 2-tuple of formal parameters; you can recognize the type of
`batch_loss` as `(<MODEL_TYPE,BATCH_TYPE> -> float32)`. `batch_loss` as `(<MODEL_TYPE,BATCH_TYPE> -> float32)`.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(batch_loss.type_signature) str(batch_loss.type_signature)
``` ```
%% Output %% Output
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>> -> float32)' '(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>> -> float32)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
As a sanity check, let's construct an initial model filled with zeros and As a sanity check, let's construct an initial model filled with zeros and
compute the loss over the batch of data we visualized above. compute the loss over the batch of data we visualized above.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
initial_model = collections.OrderedDict( initial_model = collections.OrderedDict(
weights=np.zeros([784, 10], dtype=np.float32), weights=np.zeros([784, 10], dtype=np.float32),
bias=np.zeros([10], dtype=np.float32)) bias=np.zeros([10], dtype=np.float32))
sample_batch = federated_train_data[5][-1] sample_batch = federated_train_data[5][-1]
batch_loss(initial_model, sample_batch) batch_loss(initial_model, sample_batch)
``` ```
%% Output %% Output
2.3025854 2.3025854
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Note that we feed the TFF computation with the initial model defined as a Note that we feed the TFF computation with the initial model defined as a
`dict`, even though the body of the Python function that defines it consumes `dict`, even though the body of the Python function that defines it consumes
model parameters as `model.weight` and `model.bias`. The arguments of the call model parameters as `model.weight` and `model.bias`. The arguments of the call
to `batch_loss` aren't simply passed to the body of that function. to `batch_loss` aren't simply passed to the body of that function.
What happens when we invoke `batch_loss`? What happens when we invoke `batch_loss`?
The Python body of `batch_loss` has already been traced and serialized in the above cell where it was defined. TFF acts as the caller to `batch_loss` The Python body of `batch_loss` has already been traced and serialized in the above cell where it was defined. TFF acts as the caller to `batch_loss`
at the computation definition time, and as the target of invocation at the time at the computation definition time, and as the target of invocation at the time
`batch_loss` is invoked. In both roles, TFF serves as the bridge between TFF's `batch_loss` is invoked. In both roles, TFF serves as the bridge between TFF's
abstract type system and Python representation types. At the invocation time, abstract type system and Python representation types. At the invocation time,
TFF will accept most standard Python container types (`dict`, `list`, `tuple`, TFF will accept most standard Python container types (`dict`, `list`, `tuple`,
`collections.namedtuple`, etc.) as concrete representations of abstract TFF `collections.namedtuple`, etc.) as concrete representations of abstract TFF
tuples. Also, although as noted above, TFF computations formally only accept a tuples. Also, although as noted above, TFF computations formally only accept a
single parameter, you can use the familiar Python call syntax with positional single parameter, you can use the familiar Python call syntax with positional
and/or keyword arguments in case where the type of the parameter is a tuple - it and/or keyword arguments in case where the type of the parameter is a tuple - it
works as expected. works as expected.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Gradient descent on a single batch ### Gradient descent on a single batch
Now, let's define a computation that uses this loss function to perform a single Now, let's define a computation that uses this loss function to perform a single
step of gradient descent. Note how in defining this function, we use step of gradient descent. Note how in defining this function, we use
`batch_loss` as a subcomponent. You can invoke a computation constructed with `batch_loss` as a subcomponent. You can invoke a computation constructed with
`tff.tf_computation` inside the body of another computation, though typically `tff.tf_computation` inside the body of another computation, though typically
this is not necessary - as noted above, because serialization looses some this is not necessary - as noted above, because serialization looses some
debugging information, it is often preferable for more complex computations to debugging information, it is often preferable for more complex computations to
write and test all the TensorFlow without the `tff.tf_computation` decorator. write and test all the TensorFlow without the `tff.tf_computation` decorator.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32) @tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def batch_train(initial_model, batch, learning_rate): def batch_train(initial_model, batch, learning_rate):
# Define a group of model variables and set them to `initial_model`. Must # Define a group of model variables and set them to `initial_model`. Must
# be defined outside the @tf.function. # be defined outside the @tf.function.
model_vars = collections.OrderedDict([ model_vars = collections.OrderedDict([
(name, tf.Variable(name=name, initial_value=value)) (name, tf.Variable(name=name, initial_value=value))
for name, value in initial_model.items() for name, value in initial_model.items()
]) ])
optimizer = tf.keras.optimizers.SGD(learning_rate) optimizer = tf.keras.optimizers.SGD(learning_rate)
@tf.function @tf.function
def _train_on_batch(model_vars, batch): def _train_on_batch(model_vars, batch):
# Perform one step of gradient descent using loss from `batch_loss`. # Perform one step of gradient descent using loss from `batch_loss`.
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
loss = forward_pass(model_vars, batch) loss = forward_pass(model_vars, batch)
grads = tape.gradient(loss, model_vars) grads = tape.gradient(loss, model_vars)
optimizer.apply_gradients( optimizer.apply_gradients(
zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars))) zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))
return model_vars return model_vars
return _train_on_batch(model_vars, batch) return _train_on_batch(model_vars, batch)
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(batch_train.type_signature) str(batch_train.type_signature)
``` ```
%% Output %% Output
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>,float32> -> <weights=float32[784,10],bias=float32[10]>)' '(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>,float32> -> <weights=float32[784,10],bias=float32[10]>)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
When you invoke a Python function decorated with `tff.tf_computation` within the When you invoke a Python function decorated with `tff.tf_computation` within the
body of another such function, the logic of the inner TFF computation is body of another such function, the logic of the inner TFF computation is
embedded (essentially, inlined) in the logic of the outer one. As noted above, embedded (essentially, inlined) in the logic of the outer one. As noted above,
if you are writing both computations, it is likely preferable to make the inner if you are writing both computations, it is likely preferable to make the inner
function (`batch_loss` in this case) a regular Python or `tf.function` rather function (`batch_loss` in this case) a regular Python or `tf.function` rather
than a `tff.tf_computation`. However, here we illustrate that calling one than a `tff.tf_computation`. However, here we illustrate that calling one
`tff.tf_computation` inside another basically works as expected. This may be `tff.tf_computation` inside another basically works as expected. This may be
necessary if, for example, you do not have the Python code defining necessary if, for example, you do not have the Python code defining
`batch_loss`, but only its serialized TFF representation. `batch_loss`, but only its serialized TFF representation.
Now, let's apply this function a few times to the initial model to see whether Now, let's apply this function a few times to the initial model to see whether
the loss decreases. the loss decreases.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
model = initial_model model = initial_model
losses = [] losses = []
for _ in range(5): for _ in range(5):
model = batch_train(model, sample_batch, 0.1) model = batch_train(model, sample_batch, 0.1)
losses.append(batch_loss(model, sample_batch)) losses.append(batch_loss(model, sample_batch))
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
losses losses
``` ```
%% Output %% Output
[0.19690028, 0.13176318, 0.101132266, 0.08273812, 0.0703014] [0.19690028, 0.13176318, 0.101132266, 0.08273812, 0.0703014]
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Gradient descent on a sequence of local data ### Gradient descent on a sequence of local data
Now, since `batch_train` appears to work, let's write a similar training Now, since `batch_train` appears to work, let's write a similar training
function `local_train` that consumes the entire sequence of all batches from one function `local_train` that consumes the entire sequence of all batches from one
user instead of just a single batch. The new computation will need to now user instead of just a single batch. The new computation will need to now
consume `tff.SequenceType(BATCH_TYPE)` instead of `BATCH_TYPE`. consume `tff.SequenceType(BATCH_TYPE)` instead of `BATCH_TYPE`.
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE) LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)
@tff.federated_computation(MODEL_TYPE, tf.float32, LOCAL_DATA_TYPE) @tff.federated_computation(MODEL_TYPE, tf.float32, LOCAL_DATA_TYPE)
def local_train(initial_model, learning_rate, all_batches): def local_train(initial_model, learning_rate, all_batches):
# Mapping function to apply to each batch. # Mapping function to apply to each batch.
@tff.federated_computation(MODEL_TYPE, BATCH_TYPE) @tff.federated_computation(MODEL_TYPE, BATCH_TYPE)
def batch_fn(model, batch): def batch_fn(model, batch):
return batch_train(model, batch, learning_rate) return batch_train(model, batch, learning_rate)
return tff.sequence_reduce(all_batches, initial_model, batch_fn) return tff.sequence_reduce(all_batches, initial_model, batch_fn)
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(local_train.type_signature) str(local_train.type_signature)
``` ```
%% Output %% Output
'(<<weights=float32[784,10],bias=float32[10]>,float32,<x=float32[?,784],y=int32[?]>*> -> <weights=float32[784,10],bias=float32[10]>)' '(<<weights=float32[784,10],bias=float32[10]>,float32,<x=float32[?,784],y=int32[?]>*> -> <weights=float32[784,10],bias=float32[10]>)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
There are quite a few details buried in this short section of code, let's go There are quite a few details buried in this short section of code, let's go
over them one by one. over them one by one.
First, while we could have implemented this logic entirely in TensorFlow, First, while we could have implemented this logic entirely in TensorFlow,
relying on `tf.data.Dataset.reduce` to process the sequence similarly to how relying on `tf.data.Dataset.reduce` to process the sequence similarly to how
we've done it earlier, we've opted this time to express the logic in the glue we've done it earlier, we've opted this time to express the logic in the glue
language, as a `tff.federated_computation`. We've used the federated operator language, as a `tff.federated_computation`. We've used the federated operator
`tff.sequence_reduce` to perform the reduction. `tff.sequence_reduce` to perform the reduction.
The operator `tff.sequence_reduce` is used similarly to The operator `tff.sequence_reduce` is used similarly to
`tf.data.Dataset.reduce`. You can think of it as essentially the same as `tf.data.Dataset.reduce`. You can think of it as essentially the same as
`tf.data.Dataset.reduce`, but for use inside federated computations, which as `tf.data.Dataset.reduce`, but for use inside federated computations, which as
you may remember, cannot contain TensorFlow code. It is a template operator with you may remember, cannot contain TensorFlow code. It is a template operator with
a formal parameter 3-tuple that consists of a *sequence* of `T`-typed elements, a formal parameter 3-tuple that consists of a *sequence* of `T`-typed elements,
the initial state of the reduction (we'll refer to it abstractly as *zero*) of the initial state of the reduction (we'll refer to it abstractly as *zero*) of
some type `U`, and the *reduction operator* of type `(<U,T> -> U)` that alters the some type `U`, and the *reduction operator* of type `(<U,T> -> U)` that alters the
state of the reduction by processing a single element. The result is the final state of the reduction by processing a single element. The result is the final
state of the reduction, after processing all elements in a sequential order. In state of the reduction, after processing all elements in a sequential order. In
our example, the state of the reduction is the model trained on a prefix of the our example, the state of the reduction is the model trained on a prefix of the
data, and the elements are data batches. data, and the elements are data batches.
Second, note that we have again used one computation (`batch_train`) as a Second, note that we have again used one computation (`batch_train`) as a
component within another (`local_train`), but not directly. We can't use it as a component within another (`local_train`), but not directly. We can't use it as a
reduction operator because it takes an additional parameter - the learning rate. reduction operator because it takes an additional parameter - the learning rate.
To resolve this, we define an embedded federated computation `batch_fn` that To resolve this, we define an embedded federated computation `batch_fn` that
binds to the `local_train`'s parameter `learning_rate` in its body. It is binds to the `local_train`'s parameter `learning_rate` in its body. It is
allowed for a child computation defined this way to capture a formal parameter allowed for a child computation defined this way to capture a formal parameter
of its parent as long as the child computation is not invoked outside the body of its parent as long as the child computation is not invoked outside the body
of its parent. You can think of this pattern as an equivalent of of its parent. You can think of this pattern as an equivalent of
`functools.partial` in Python. `functools.partial` in Python.
The practical implication of capturing `learning_rate` this way is, of course, The practical implication of capturing `learning_rate` this way is, of course,
that the same learning rate value is used across all batches. that the same learning rate value is used across all batches.
Now, let's try the newly defined local training function on the entire sequence Now, let's try the newly defined local training function on the entire sequence
of data from the same user who contributed the sample batch (digit `5`). of data from the same user who contributed the sample batch (digit `5`).
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5]) locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5])
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Did it work? To answer this question, we need to implement evaluation. Did it work? To answer this question, we need to implement evaluation.
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Local evaluation ### Local evaluation
Here's one way to implement local evaluation by adding up the losses across all data Here's one way to implement local evaluation by adding up the losses across all data
batches (we could have just as well computed the average; we'll leave it as an batches (we could have just as well computed the average; we'll leave it as an
exercise for the reader). exercise for the reader).
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE) @tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)
def local_eval(model, all_batches): def local_eval(model, all_batches):
# TODO(b/120157713): Replace with `tff.sequence_average()` once implemented. # TODO(b/120157713): Replace with `tff.sequence_average()` once implemented.
return tff.sequence_sum( return tff.sequence_sum(
tff.sequence_map( tff.sequence_map(
tff.federated_computation(lambda b: batch_loss(model, b), BATCH_TYPE), tff.federated_computation(lambda b: batch_loss(model, b), BATCH_TYPE),
all_batches)) all_batches))
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` ```
str(local_eval.type_signature) str(local_eval.type_signature)
``` ```
%% Output %% Output
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>*> -> float32)' '(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>*> -> float32)'
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Again, there are a few new elements illustrated by this code, let's go over them Again, there are a few new elements illustrated by this code, let's go over them
one by one. one by one.
First, we have used two new federated operators for processing sequences: First, we have used two new federated operators for processing sequences:
`tff.sequence_map` that takes a *mapping function* `T->U` and a *sequence* of `tff.sequence_map` that takes a *mapping function* `T->U` and a *sequence* of
`T`, and emits a sequence of `U` obtained by applying the mapping function `T`, and emits a sequence of `U` obtained by applying the mapping function
pointwise, and `tff.sequence_sum` that just adds all the elements. Here, we map pointwise, and `tff.sequence_sum` that just adds all the elements. Here, we map
each data batch to a loss value, and then add the resulting loss values to each data batch to a loss value, and then add the resulting loss values to
compute the total loss. compute the total loss.
Note that we could have again used `tff.sequence_reduce`, but this wouldn't be Note that we could have again used `tff.sequence_reduce`, but this wouldn't be
the best choice - the reduction process is, by definition, sequential, whereas the best choice - the reduction process is, by definition, sequential, whereas
the mapping and sum can be computed in parallel. When given a choice, it's best the mapping and sum can be computed in parallel. When given a choice, it's best
to stick with operators that don't constrain implementation choices, so that to stick with operators that don't constrain implementation choices, so that
when our TFF computation is compiled in the future to be deployed to a specific when our TFF computation is compiled in the future to be deployed to a specific
environment, one can take full advantage of all potential opportunities for a environment, one can take full advantage of all potential opportunities for a
faster, more scalable, more resource-efficient execution. faster, more scalable, more resource-efficient execution.
Second, note that just as in `local_train`, the component function we need Second, note that just as in `local_train`, the component function we need
(`batch_loss`) takes more parameters than what the federated operator (`batch_loss`) takes more parameters than what the federated operator
(`tff.sequence_map`) expects, so we again define a partial, this time inline by (`tff.sequence_map`) expects, so we again define a partial, this time inline by
directly wrapping a `lambda` as a `tff.federated_computation`. Using wrappers directly wrapping a `lambda` as a `tff.federated_computation`. Using wrappers
inline with a function as an argument is the recommended way to use inline with a function as an argument is the recommended way to use
`tff.tf_computation` to embed TensorFlow logic in TFF. `tff.tf_computation` to embed TensorFlow logic in TFF.