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

Creating TFF release 0.13.0 

PiperOrigin-RevId: 300591979
上级 a8ea2945
Hide whitespace changes
Inline Side-by-side
README.md
浏览文件 @ 7c0121a5
......@@ -91,6 +91,7 @@ versa.
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.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/)
......
RELEASE.md
浏览文件 @ 7c0121a5
# 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
## Major Features and Improvements
......
docs/tutorials/custom_federated_algorithms_1.ipynb
浏览文件 @ 7c0121a5
%% Cell type:markdown id: tags:
##### Copyright 2019 The TensorFlow Authors.
%% Cell type:code id: tags:
```
#@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 obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
```
%% Cell type:markdown id: tags:
# Custom Federated Algorithms, Part 1: Introduction to the Federated Core
%% Cell type:markdown id: tags:
<table class="tfo-notebook-buttons" align="left">
<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>
</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>
<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>
</table>
%% Cell type:markdown id: tags:
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)
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
[Federated Learning (FL)](../federated_learning.md) layer.
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
simple example with a distributed array of temperature sensors. In
[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
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)
of federated averaging in `tff.learning`.
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
abstractions we offer are quite generic, and could be used, e.g., to implement
analytics and other custom types of computations over distributed data.
Although this tutorial is designed to be self-contained, we encourage you to
first read tutorials on
[image classification](federated_learning_for_image_classification.ipynb) and
[text generation](federated_learning_for_text_generation.ipynb) for a
higher-level and more gentle introduction to the TensorFlow Federated framework
and the [Federated Learning](../federated_learning.md) APIs (`tff.learning`), as
it will help you put the concepts we describe here in context.
%% Cell type:markdown id: tags:
## Intended Uses
In a nutshell, Federated Core (FC) is a development environment that makes it
possible to compactly express program logic that combines TensorFlow code with
distributed communication operators, such as those that are used in
[Federated Averaging](https://arxiv.org/abs/1602.05629) - computing
distributed sums, averages, and other types of distributed aggregations over a
set of client devices in the system, broadcasting models and parameters to those
devices, etc.
You may be aware of
[`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
framework differ? Both frameworks attempt at making TensorFlow computations
distributed, after all.
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
code with minimal changes to enable distributed training*, and much focus is on
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
practitioners explicit control over the specific patterns of distributed
communication they will use in their systems. The focus in FC is on providing a
flexible and extensible language for expressing distributed data flow
algorithms, rather than a concrete set of implemented distributed training
capabilities.
One of the primary target audiences for TFF's FC API is researchers and
practitioners who might want to experiment with new federated learning
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,
yet without getting bogged down by system implementation details. The level of
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
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
exchanges.
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
centralized location may not be a viable option. This has implication on the
implementation of machine learning algorithms that require an increased degree
of explicit control, as compared to scenarios in which all data can be
accumulated in a centralized location at a data center.
%% Cell type:markdown id: tags:
## Before we start
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,
please refer to the [Installation](../install.md) guide for instructions.
%% Cell type:code id: tags:
```
#@test {"skip": true}
!pip install --quiet --upgrade tensorflow_federated
# Note: Jupyter requires a patch to asyncio.
!pip install --quiet --upgrade nest_asyncio
import nest_asyncio
nest_asyncio.apply()
```
%% Cell type:code id: tags:
```
import collections
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
tf.compat.v1.enable_v2_behavior()
```
%% Cell type:code id: tags:
```
@tff.federated_computation
def hello_world():
return 'Hello, World!'
hello_world()
```
%% Output
b'Hello, World!'
%% Cell type:markdown id: tags:
## Federated data
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
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,
applications running on mobile devices may collect data and store it locally,
without uploading to a centralized location. Or, an array of distributed sensors
may collect and store temperature readings at their locations.
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.,
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
values*, or as *values of federated types*.
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
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*
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
of this federated type.
%% Cell type:code id: tags:
```
federated_float_on_clients = tff.FederatedType(tf.float32, tff.CLIENTS)
```
%% Cell type:markdown id: tags:
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,
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
the group `G` of devices hosting a federated value as the value's *placement*.
Thus, `tff.CLIENTS` is an example of a placement.
%% Cell type:code id: tags:
```
str(federated_float_on_clients.member)
```
%% Output
'float32'
%% Cell type:code id: tags:
```
str(federated_float_on_clients.placement)
```
%% Output
'CLIENTS'
%% Cell type:markdown id: tags:
A federated type with member constituents `T` and placement `G` can be
represented compactly as `{T}@G`, as shown below.
%% Cell type:code id: tags:
```
str(federated_float_on_clients)
```
%% Output
'{float32}@CLIENTS'
%% Cell type:markdown id: tags:
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
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
`T`-typed items that together constitute the federated value.
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
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
operators* that abstractly represent various kinds of distributed communication
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.
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
known to be all equal. This is controlled by the third, optional `all_equal`
parameter in the `tff.FederatedType` constructor (defaulting to `False`).
%% Cell type:code id: tags:
```
federated_float_on_clients.all_equal
```
%% Output
False
%% Cell type:markdown id: tags:
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
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).
%% Cell type:code id: tags:
```
str(tff.FederatedType(tf.float32, tff.CLIENTS, all_equal=True))
```
%% Output
'float32@CLIENTS'
%% Cell type:markdown id: tags:
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.)
that has been broadcasted by a server to a group of devices that participate in
federated training.
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
they can be personalized for each user.
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
(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
tuples.
%% Cell type:code id: tags:
```
simple_regression_model_type = (
tff.NamedTupleType([('a', tf.float32), ('b', tf.float32)]))
str(simple_regression_model_type)
```
%% Output
'<a=float32,b=float32>'
%% Cell type:markdown id: tags:
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
general `tff.TensorType(dtype=tf.float32, shape=[])`.
When this model is broadcasted to clients, the type of the resulting federated
value can be represented as shown below.
%% Cell type:code id: tags:
```
str(tff.FederatedType(
simple_regression_model_type, tff.CLIENTS, all_equal=True))
```
%% Output
'<a=float32,b=float32>@CLIENTS'
%% Cell type:markdown id: tags:
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
refer to a federated value in which member constituents are *XYZ*-like. Thus, we
will talk about things like *federated tuples*, *federated sequences*,
*federated models*, and so on.
Now, coming back to `float32@CLIENTS` - while it appears replicated across
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
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`.
Given the isomorphism between `T` and `T@G`, you may wonder what purpose, if
any, the latter types might serve. Read on.
%% Cell type:markdown id: tags:
## Placements
### Design Overview
In the preceding section, we've introduced the concept of *placements* - groups
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
placement.
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
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
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
system, potentially including mobile or embedded devices running Android. Each
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
centralized coordinator, an intermediate layer in a multi-tier architecture,
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
materialize.
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
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
the mechanisms designed to support this.
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
those operations might *run*, TFF focuses on *data*, where that data
*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
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
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
`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
formal privacy guarantees for sensitive on-device data.
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
data (the placements), the programmer will never deal with the raw data or
identities of the *individual* participants.
(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`
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
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
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
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
client identity anywhere in the Federated Core API, the underlying set of
architectural abstractions, or the core runtime infrastructure we provide to
support simulations. All the computation logic you write will be expressed as
operations on the entire client group.
Recall here what we mentioned earlier about values of federated types being
unlike Python `dict`, in that one cannot simply enumerate their member
constituents. Think of values that your TFF program logic manipulates as being
associated with placements (groups), rather than with individual participants.
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
`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
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
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`.
### Specifying Placements
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
naturally modeled as client-server architectures, with multiple *client* devices
(mobile phones, embedded devices, distributed databases, sensors, etc.)
orchestrated by a single centralized *server* coordinator. TFF is designed to
also support custom placements, multiple client groups, multi-tiered and other,
more general distributed architectures, but discussing them is outside the scope
of this tutorial.
TFF doesn't prescribe what either the `tff.CLIENTS` or the `tff.SERVER` actually
represent.
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
fault-tolerant cluster running state machine replication - we do not make any
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
only a single item of data at the server.
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
*population*, but e.g., in
[production implementations of Federated Averaging](https://arxiv.org/abs/1602.05629),
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
concrete meaning when a computation in which they appear is deployed for
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
clients is determined by the federated data supplied as input.
%% Cell type:markdown id: tags:
## Federated computations
### Declaring federated computations
TFF is designed as a strongly-typed functional programming environment that
supports modular development.
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
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.
%% Cell type:code id: tags:
```
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(sensor_readings):
return tff.federated_mean(sensor_readings)
```
%% Cell type:markdown id: tags:
Looking at the above code, at this point you might be asking - aren't there
already decorator constructs to define composable units such as
[`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
different?
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
distributed system in an internal platform-independent *glue* language. At this
point, this will undoubtedly sound cryptic, but please bear this intuitive
interpretation of a federated computation as an abstract specification of a
distributed system in mind. We'll explain it in a minute.
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
signatures. You can print the type signature of a computation by querying its
`type_signature` property, as shown below.
%% Cell type:code id: tags:
```
str(get_average_temperature.type_signature)
```
%% Output
'({float32}@CLIENTS -> float32@SERVER)'
%% Cell type:markdown id: tags:
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
server.
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
`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
just said about federated computations as representing abstract specifications
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
aggregate results emerge at the server.
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
reflects the idea that computations might be triggered by *queries* that
originate and terminate on the server.
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
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
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
computations, it's best to always default to the global network-wide
perspective, rather than the perspective of a single centralized coordinator.
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
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 -
it's determined automatically.
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
understanding, debugging, and formally verifying properties of your code easier.
In some cases, explicitly specifying types is a requirement (e.g., polymorphic
computations are currently not directly executable).
### Executing federated computations
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
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
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.
%% Cell type:code id: tags:
```
get_average_temperature([68.5, 70.3, 69.8])
```
%% Output
69.53334
%% Cell type:markdown id: tags:
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
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.
Now, let's return to a note we made earlier about the
`tff.federated_computation` decorator emitting code in a *glue* language.
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
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
earlier, TFF computations are actually *not* Python.
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
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)
of the computation's logic for future use - whether for execution, or to be
incorporated as a sub-component into another computation.
You can verify this by adding a print statement, as follows:
%% Cell type:code id: tags:
```
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(sensor_readings):
print ('Getting traced, the argument is "{}".'.format(
type(sensor_readings).__name__))
return tff.federated_mean(sensor_readings)
```
%% Output
Getting traced, the argument is "ValueImpl".
%% Cell type:markdown id: tags:
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
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
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
platform-independent and serializable.
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,
are compiled into a portable and platform-independent serializable
representation under the hood.
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
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
TFF computation's body is executed at definition time, when the body of the
Python function decorated with `tff.federated_computation` is traced before
getting serialized. It's not retraced again at invocation time (except when the
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
representation. One reason is that ultimately, TFF computations are intended to
be deployable to real physical environments, and hosted on mobile or embedded
devices, where Python may not be available.
Another reason is that TFF computations express the global behavior of
distributed systems, as opposed to Python programs which express the local
behavior of individual participants. You can see that in the simple example
above, with the special operator `tff.federated_mean` that accepts data on
client devices, but deposits the results on the server.
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
represents a distributed system that coordinates the behavior of multiple system
participants. We will refer to such operators as *federated operators*, to
distinguish them from ordinary (local) operators in Python.
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
use of a dedicated representation.
### Composing federated computations
As noted above, federated computations and their constituents are best
understood as models of distributed systems, and you can think of composing
federated computations as composing more complex distributed systems from
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 ->
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).
The same is true of composing federated computations. The computation
`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
embedded in the body of the parent, much in the same way `tff.federated_mean`
was embedded in its own body earlier.
An important restriction to be aware of is that bodies of Python functions
decorated with `tff.federated_computation` must consist *only* of federated
operators, i.e., they cannot directly contain TensorFlow operations. For
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
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
a `tff.federated_computation`.
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
federated computations (i.e., the logic constructed from serialized bodies of
Python functions decorated with `tff.federated_computation`) is designed to
serve as a platform-independent *glue* language. This glue language is currently
used to build distributed systems from embedded sections of TensorFlow code
(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
relational database queries that might represent input pipelines, all connected
together using the same glue language (the `tff.federated_computation` blocks).
%% Cell type:markdown id: tags:
## TensorFlow logic
### Declaring TensorFlow computations
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.
In order to use such code with TFF, as noted above, it just needs to be
decorated with `tff.tf_computation`.
For example, here's how we could implement a function that takes a number and
adds `0.5` to it.
%% Cell type:code id: tags:
```
@tff.tf_computation(tf.float32)
def add_half(x):
return tf.add(x, 0.5)
```
%% Cell type:markdown id: tags:
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
such as `tf.function`. Unlike in the preceding section, here we are
dealing with an ordinary block of TensorFlow code.
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:
* In order to embed reusable building blocks implemented using TensorFlow code
in the bodies of federated computations, they need to satisfy certain
properties - such as getting traced and serialized at definition time,
having type signatures, etc. This generally requires some form of a
decorator.
In general, we recommend using TensorFlow's native mechanisms for composition,
such as `tf.function`, wherever possible, as the exact manner in
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`
we just defined can be treated by TFF just like any other TFF computation. In
particular, it has a TFF type signature.
%% Cell type:code id: tags:
```
str(add_half.type_signature)
```
%% Output
'(float32 -> float32)'
%% Cell type:markdown id: tags:
Note this type signature does not have placements. TensorFlow computations
cannot consume or return federated types.
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
`add_half` pointwise to all member constituents of a federated float on client
devices.
%% Cell type:code id: tags:
```
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def add_half_on_clients(x):
return tff.federated_map(add_half, x)
```
%% Cell type:code id: tags:
```
str(add_half_on_clients.type_signature)
```
%% Output
'({float32}@CLIENTS -> {float32}@CLIENTS)'
%% Cell type:markdown id: tags:
### Executing TensorFlow computations
Execution of computations defined with `tff.tf_computation` follows the same
rules as those we described for `tff.federated_computation`. They can be invoked
as ordinary callables in Python, as follows.
%% Cell type:code id: tags:
```
add_half_on_clients([1.0, 3.0, 2.0])
```
%% Output
[<tf.Tensor: shape=(), dtype=float32, numpy=1.5>,
<tf.Tensor: shape=(), dtype=float32, numpy=3.5>,
<tf.Tensor: shape=(), dtype=float32, numpy=2.5>]
%% Cell type:markdown id: tags:
Once again, it is worth noting that invoking the computation
`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
client perform a local action. There is no `tff.SERVER` explicitly mentioned in
this system (even if in practice, orchestrating such processing might involve
one). Think of a computation defined this way as conceptually analogous to the
`Map` stage in `MapReduce`.
Also, keep in mind that what we said in the preceding section about TFF
computations getting serialized at the definition time remains true for
`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
serialized representation.
The only difference between Python methods decorated with
`tff.federated_computation` and those decorated with `tff.tf_computation` is
that the latter are serialized as TensorFlow graphs (whereas the former are not
allowed to contain TensorFlow code directly embedded in them).
Under the hood, each method decorated with `tff.tf_computation` temporarily
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
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
serialized.
For example, the following code will fail:
%% Cell type:code id: tags:
```
try:
# Eager mode
constant_10 = tf.constant(10.)
@tff.tf_computation(tf.float32)
def add_ten(x):
return x + constant_10
except Exception as err:
print (err)
```
%% Output
Attempting to capture an EagerTensor without building a function.
%% Cell type:markdown id: tags:
The above fails because `constant_10` has already been constructed outside of
the graph that `tff.tf_computation` constructs internally in the body of
`add_ten` during the serialization process.
On the other hand, invoking python functions that modify the current graph when
called inside a `tff.tf_computation` is fine:
%% Cell type:code id: tags:
```
def get_constant_10():
return tf.constant(10.)
@tff.tf_computation(tf.float32)
def add_ten(x):
return x + get_constant_10()
add_ten(5.0)
```
%% Output
15.0
%% Cell type:markdown id: tags:
Note that the serialization mechanisms in TensorFlow are evolving, and we expect
the details of how TFF serializes computations to evolve as well.
### Working with `tf.data.Dataset`s
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
your code. Parameters to be represented in TensorFlow as data sets need to be
declared using the `tff.SequenceType` constructor.
For example, the type specification `tff.SequenceType(tf.float32)` defines an
abstract sequence of float elements in TFF. Sequences can contain either
tensors, or complex nested structures (we'll see examples of those later). The
concise representation of a sequence of `T`-typed items is `T*`.
%% Cell type:code id: tags:
```
float32_sequence = tff.SequenceType(tf.float32)
str(float32_sequence)
```
%% Output
'float32*'
%% Cell type:markdown id: tags:
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
in TensorFlow that calculates the average of temperatures in a single local data
set using the `tf.data.Dataset.reduce` operator.
%% Cell type:code id: tags:
```
@tff.tf_computation(tff.SequenceType(tf.float32))
def get_local_temperature_average(local_temperatures):
sum_and_count = (
local_temperatures.reduce((0.0, 0), lambda x, y: (x[0] + y, x[1] + 1)))
return sum_and_count[0] / tf.cast(sum_and_count[1], tf.float32)
```
%% Cell type:code id: tags:
```
str(get_local_temperature_average.type_signature)
```
%% Output
'(float32* -> float32)'
%% Cell type:markdown id: tags:
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
`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
support for data sets in TensorFlow evolves).
You can easily verify this as follows.
%% Cell type:code id: tags:
```
@tff.tf_computation(tff.SequenceType(tf.int32))
def foo(x):
return x.reduce(np.int32(0), lambda x, y: x + y)
foo([1, 2, 3])
```
%% Output
6
%% Cell type:markdown id: tags:
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
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
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
refer to the documentation pages for details).
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
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).
%% Cell type:code id: tags:
```
get_local_temperature_average([68.5, 70.3, 69.8])
```
%% Output
69.53333
%% Cell type:markdown id: tags:
Like all other TFF types, sequences like those defined above can use the
`tff.NamedTupleType` constructor to define nested structures. For example,
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
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`.
%% Cell type:code id: tags:
```
@tff.tf_computation(tff.SequenceType(collections.OrderedDict([('A', tf.int32), ('B', tf.int32)])))
def foo(ds):
print('element_structure = {}'.format(ds.element_spec))
return ds.reduce(np.int32(0), lambda total, x: total + x['A'] * x['B'])
```
%% Output
element_structure = OrderedDict([('A', TensorSpec(shape=(), dtype=tf.int32, name=None)), ('B', TensorSpec(shape=(), dtype=tf.int32, name=None))])
%% Cell type:code id: tags:
```
str(foo.type_signature)
```
%% Output
'(<A=int32,B=int32>* -> int32)'
%% Cell type:code id: tags:
```
foo([{'A': 2, 'B': 3}, {'A': 4, 'B': 5}])
```
%% Output
26
%% Cell type:markdown id: tags:
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
in this tutorial.
## Putting it all together
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
temperature readings. We can compute the global temperature average by averaging
the sensors' local averages as follows.
%% Cell type:code id: tags:
```
@tff.federated_computation(
tff.FederatedType(tff.SequenceType(tf.float32), tff.CLIENTS))
def get_global_temperature_average(sensor_readings):
return tff.federated_mean(
tff.federated_map(get_local_temperature_average, sensor_readings))
```
%% Cell type:markdown id: tags:
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
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
accepts the weight as an optional second argument (expected to be a federated
float).
Also note that the input to `get_global_temperature_average` now becomes a
*federated float sequence*. Federated sequences is how we will typically represent
on-device data in federated learning, with sequence elements typically
representing data batches (you will see examples of this shortly).
%% Cell type:code id: tags:
```
str(get_global_temperature_average.type_signature)
```
%% Output
'({float32*}@CLIENTS -> float32@SERVER)'
%% Cell type:markdown id: tags:
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
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.
%% Cell type:code id: tags:
```
get_global_temperature_average([[68.0, 70.0], [71.0], [68.0, 72.0, 70.0]])
```
%% Output
70.0
%% Cell type:markdown id: tags:
This concludes the first part of the tutorial... we encourage you to continue on
to the [second part](custom_federated_algorithms_2.ipynb).
......
docs/tutorials/custom_federated_algorithms_2.ipynb
浏览文件 @ 7c0121a5
%% Cell type:markdown id: tags:
##### Copyright 2019 The TensorFlow Authors.
%% Cell type:code id: tags:
```
#@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 obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
```
%% Cell type:markdown id: tags:
# Custom Federated Algorithms, Part 2: Implementing Federated Averaging
%% Cell type:markdown id: tags:
<table class="tfo-notebook-buttons" align="left">
<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>
</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>
<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>
</table>
%% Cell type:markdown id: tags:
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
[Federated Core (FC)](../federated_core.md), which serves as a foundation for
the [Federated Learning (FL)](../federated_learning.md) layer (`tff.learning`).
We encourage you to first read the
[first part of this series](custom_federated_algorithms_1.ipynb), which
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
to implement a simple version of federated training and evaluation algorithms.
We encourage you to review the
[image classification](federated_learning_for_image_classification.ipynb) and
[text generation](federated_learning_for_text_generation.ipynb) tutorials for a
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.
%% Cell type:markdown id: tags:
## Before we start
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
[Installation](../install.md) guide for instructions.
%% Cell type:code id: tags:
```
#@test {"skip": true}
!pip install --quiet --upgrade tensorflow_federated
```
%% Cell type:code id: tags:
```
import collections
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
tf.compat.v1.enable_v2_behavior()
# TODO(b/148678573,b/148685415): must use the ReferenceExecutor because it
# supports unbounded references and tff.sequence_* intrinsics.
tff.framework.set_default_executor(tff.framework.ReferenceExecutor())
```
%% Cell type:code id: tags:
```
@tff.federated_computation
def hello_world():
return 'Hello, World!'
hello_world()
```
%% Output
'Hello, World!'
%% Cell type:markdown id: tags:
## Implementing Federated Averaging
As in
[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
tutorial, we are going to bypass the Keras API and `tff.simulation`, write raw
model code, and construct a federated data set from scratch.
%% Cell type:markdown id: tags:
### Preparing federated data sets
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
recognize a different digit. This is about as
non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables)
as it gets.
First, let's load the standard MNIST data:
%% Cell type:code id: tags:
```
mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
```
%% Cell type:code id: tags:
```
[(x.dtype, x.shape) for x in mnist_train]
```
%% Output
[(dtype('uint8'), (60000, 28, 28)), (dtype('uint8'), (60000,))]
%% Cell type:markdown id: tags:
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
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
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
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
pixels in it into the `0..1` range, so that we don't have to clutter the model
logic with data conversions.
%% Cell type:code id: tags:
```
NUM_EXAMPLES_PER_USER = 1000
BATCH_SIZE = 100
def get_data_for_digit(source, digit):
output_sequence = []
all_samples = [i for i, d in enumerate(source[1]) if d == digit]
for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):
batch_samples = all_samples[i:i + BATCH_SIZE]
output_sequence.append({
'x':
np.array([source[0][i].flatten() / 255.0 for i in batch_samples],
dtype=np.float32),
'y':
np.array([source[1][i] for i in batch_samples], dtype=np.int32)
})
return output_sequence
federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)]
federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]
```
%% Cell type:markdown id: tags:
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`).
%% Cell type:code id: tags:
```
federated_train_data[5][-1]['y']
```
%% Output
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], dtype=int32)
%% Cell type:markdown id: tags:
Just to be sure, let's also look at the image corresponding to the last element of that batch.
%% Cell type:code id: tags:
```
from matplotlib import pyplot as plt
plt.imshow(federated_train_data[5][-1]['x'][-1].reshape(28, 28), cmap='gray')
plt.grid(False)
plt.show()
```
%% Output
%% Cell type:markdown id: tags:
### On combining TensorFlow and TFF
In this tutorial, for compactness we immediately decorate functions that
introduce TensorFlow logic with `tff.tf_computation`. However, for more complex
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
then re-imported necessarily loses some metadata and limits interactivity,
making debugging even more of a challenge.
Therefore, **we strongly recommend writing complex TF logic as stand-alone
Python functions** (that is, without `tff.tf_computation` decoration). This way
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).
%% Cell type:markdown id: tags:
### Defining a loss function
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
size of data batches may vary, we set the batch dimension to `None` to indicate
that the size of this dimension is unknown.
%% Cell type:code id: tags:
```
BATCH_SPEC = collections.OrderedDict(
x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),
y=tf.TensorSpec(shape=[None], dtype=tf.int32))
BATCH_TYPE = tff.to_type(BATCH_SPEC)
str(BATCH_TYPE)
```
%% Output
'<x=float32[?,784],y=int32[?]>'
%% Cell type:markdown id: tags:
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
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`
defined above represents an abstract TFF type specification. It is important to
distinguish this *abstract* TFF type from concrete Python *representation*
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,
TFF has a single abstract type constructor `tff.NamedTupleType` for tuple-like
containers, with elements that can be individually named or left unnamed. This
type is also used to model formal parameters of computations, as TFF
computations can formally only declare one parameter and one result - you will
see examples of this shortly.
Let's now define the TFF type of model parameters, again as a TFF named tuple of
*weights* and *bias*.
%% Cell type:code id: tags:
```
MODEL_SPEC = collections.OrderedDict(
weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),
bias=tf.TensorSpec(shape=[10], dtype=tf.float32))
MODEL_TYPE = tff.to_type(MODEL_SPEC)
print(MODEL_TYPE)
```
%% Output
<weights=float32[784,10],bias=float32[10]>
%% 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.
%% Cell type:code id: tags:
```
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE)
@tf.function
def batch_loss(model, batch):
predicted_y = tf.nn.softmax(
tf.matmul(batch['x'], model['weights']) + model['bias'])
return -tf.reduce_mean(
tf.reduce_sum(
tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))
```
%% Cell type:markdown id: tags:
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
together into a 2-tuple of formal parameters; you can recognize the type of
`batch_loss` as `(<MODEL_TYPE,BATCH_TYPE> -> float32)`.
%% Cell type:code id: tags:
```
str(batch_loss.type_signature)
```
%% Output
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>> -> float32)'
%% Cell type:markdown id: tags:
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.
%% Cell type:code id: tags:
```
initial_model = collections.OrderedDict(
weights=np.zeros([784, 10], dtype=np.float32),
bias=np.zeros([10], dtype=np.float32))
sample_batch = federated_train_data[5][-1]
batch_loss(initial_model, sample_batch)
```
%% Output
2.3025854
%% Cell type:markdown id: tags:
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
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.
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`
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
abstract type system and Python representation types. At the invocation time,
TFF will accept most standard Python container types (`dict`, `list`, `tuple`,
`collections.namedtuple`, etc.) as concrete representations of abstract TFF
tuples. Also, although as noted above, TFF computations formally only accept a
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
works as expected.
%% Cell type:markdown id: tags:
### Gradient descent on a single batch
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
`batch_loss` as a subcomponent. You can invoke a computation constructed with
`tff.tf_computation` inside the body of another computation, though typically
this is not necessary - as noted above, because serialization looses some
debugging information, it is often preferable for more complex computations to
write and test all the TensorFlow without the `tff.tf_computation` decorator.
%% Cell type:code id: tags:
```
@tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def batch_train(initial_model, batch, learning_rate):
# Define a group of model variables and set them to `initial_model`. Must
# be defined outside the @tf.function.
model_vars = collections.OrderedDict([
(name, tf.Variable(name=name, initial_value=value))
for name, value in initial_model.items()
])
optimizer = tf.keras.optimizers.SGD(learning_rate)
@tf.function
def _train_on_batch(model_vars, batch):
# Perform one step of gradient descent using loss from `batch_loss`.
with tf.GradientTape() as tape:
loss = forward_pass(model_vars, batch)
grads = tape.gradient(loss, model_vars)
optimizer.apply_gradients(
zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))
return model_vars
return _train_on_batch(model_vars, batch)
```
%% Cell type:code id: tags:
```
str(batch_train.type_signature)
```
%% Output
'(<<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:
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
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
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
`tff.tf_computation` inside another basically works as expected. This may be
necessary if, for example, you do not have the Python code defining
`batch_loss`, but only its serialized TFF representation.
Now, let's apply this function a few times to the initial model to see whether
the loss decreases.
%% Cell type:code id: tags:
```
model = initial_model
losses = []
for _ in range(5):
model = batch_train(model, sample_batch, 0.1)
losses.append(batch_loss(model, sample_batch))
```
%% Cell type:code id: tags:
```
losses
```
%% Output
[0.19690028, 0.13176318, 0.101132266, 0.08273812, 0.0703014]
%% Cell type:markdown id: tags:
### Gradient descent on a sequence of local data
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
user instead of just a single batch. The new computation will need to now
consume `tff.SequenceType(BATCH_TYPE)` instead of `BATCH_TYPE`.
%% Cell type:code id: tags:
```
LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)
@tff.federated_computation(MODEL_TYPE, tf.float32, LOCAL_DATA_TYPE)
def local_train(initial_model, learning_rate, all_batches):
# Mapping function to apply to each batch.
@tff.federated_computation(MODEL_TYPE, BATCH_TYPE)
def batch_fn(model, batch):
return batch_train(model, batch, learning_rate)
return tff.sequence_reduce(all_batches, initial_model, batch_fn)
```
%% Cell type:code id: tags:
```
str(local_train.type_signature)
```
%% Output
'(<<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:
There are quite a few details buried in this short section of code, let's go
over them one by one.
First, while we could have implemented this logic entirely in TensorFlow,
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
language, as a `tff.federated_computation`. We've used the federated operator
`tff.sequence_reduce` to perform the reduction.
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`, but for use inside federated computations, which as
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,
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
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
our example, the state of the reduction is the model trained on a prefix of the
data, and the elements are data batches.
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
reduction operator because it takes an additional parameter - the learning rate.
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
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. You can think of this pattern as an equivalent of
`functools.partial` in Python.
The practical implication of capturing `learning_rate` this way is, of course,
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
of data from the same user who contributed the sample batch (digit `5`).
%% Cell type:code id: tags:
```
locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5])
```
%% Cell type:markdown id: tags:
Did it work? To answer this question, we need to implement evaluation.
%% Cell type:markdown id: tags:
### Local evaluation
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
exercise for the reader).
%% Cell type:code id: tags:
```
@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)
def local_eval(model, all_batches):
# TODO(b/120157713): Replace with `tff.sequence_average()` once implemented.
return tff.sequence_sum(
tff.sequence_map(
tff.federated_computation(lambda b: batch_loss(model, b), BATCH_TYPE),
all_batches))
```
%% Cell type:code id: tags:
```
str(local_eval.type_signature)
```
%% Output
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>*> -> float32)'
%% Cell type:markdown id: tags:
Again, there are a few new elements illustrated by this code, let's go over them
one by one.
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
`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
each data batch to a loss value, and then add the resulting loss values to
compute the total loss.
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 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
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
faster, more scalable, more resource-efficient execution.
Second, note that just as in `local_train`, the component function we need
(`batch_loss`) takes more parameters than what the federated operator
(`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
inline with a function as an argument is the recommended way to use
`tff.tf_computation` to embed TensorFlow logic in TFF.
Now, let's see whether our training worked.
%% Cell type:code id: tags:
```
print('initial_model loss =', local_eval(initial_model,
federated_train_data[5]))
print('locally_trained_model loss =',
local_eval(locally_trained_model, federated_train_data[5]))
```
%% Output
initial_model loss = 23.025854
locally_trained_model loss = 0.4348469
%% Cell type:markdown id: tags:
Indeed, the loss decreased. But what happens if we evaluated it on another
user's data?
%% Cell type:code id: tags:
```
print('initial_model loss =', local_eval(initial_model,
federated_train_data[0]))
print('locally_trained_model loss =',
local_eval(locally_trained_model, federated_train_data[0]))
```
%% Output
initial_model loss = 23.025854
locally_trained_model loss = 74.500755
%% Cell type:markdown id: tags:
As expected, things got worse. The model was trained to recognize `5`, and has
never seen a `0`. This brings the question - how did the local training impact
the quality of the model from the global perspective?
%% Cell type:markdown id: tags:
### Federated evaluation
This is the point in our journey where we finally circle back to federated types
and federated computations - the topic that we started with. Here's a pair of
TFF types definitions for the model that originates at the server, and the data
that remains on the clients.
%% Cell type:code id: tags:
```
SERVER_MODEL_TYPE = tff.FederatedType(MODEL_TYPE, tff.SERVER)
CLIENT_DATA_TYPE = tff.FederatedType(LOCAL_DATA_TYPE, tff.CLIENTS)
```
%% Cell type:markdown id: tags:
With all the definitions introduced so far, expressing federated evaluation in
TFF is a one-liner - we distribute the model to clients, let each client invoke
local evaluation on its local portion of data, and then average out the loss.
Here's one way to write this.
%% Cell type:code id: tags:
```
@tff.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)
def federated_eval(model, data):
return tff.federated_mean(
tff.federated_map(local_eval, [tff.federated_broadcast(model), data]))
```
%% Cell type:markdown id: tags:
We've already seen examples of `tff.federated_mean` and `tff.federated_map`
in simpler scenarios, and at the intuitive level, they work as expected, but
there's more in this section of code than meets the eye, so let's go over it
carefully.
First, let's break down the *let each client invoke local evaluation on its
local portion of data* part. As you may recall from the preceding sections,
`local_eval` has a type signature of the form `(<MODEL_TYPE, LOCAL_DATA_TYPE> ->
float32)`.
The federated operator `tff.federated_map` is a template that accepts as a
parameter a 2-tuple that consists of the *mapping function* of some type `T->U`
and a federated value of type `{T}@CLIENTS` (i.e., with member constituents of
the same type as the parameter of the mapping function), and returns a result of
type `{U}@CLIENTS`.
Since we're feeding `local_eval` as a mapping function to apply on a per-client
basis, the second argument should be of a federated type `{<MODEL_TYPE,
LOCAL_DATA_TYPE>}@CLIENTS`, i.e., in the nomenclature of the preceding sections,
it should be a federated tuple. Each client should hold a full set of arguments
for `local_eval` as a member consituent. Instead, we're feeding it a 2-element
Python `list`. What's happening here?
Indeed, this is an example of an *implicit type cast* in TFF, similar to
implicit type casts you may have encountered elsewhere, e.g., when you feed an
`int` to a function that accepts a `float`. Implicit casting is used scarcily at
this point, but we plan to make it more pervasive in TFF as a way to minimize
boilerplate.
The implicit cast that's applied in this case is the equivalence between
federated tuples of the form `{<X,Y>}@Z`, and tuples of federated values
`<{X}@Z,{Y}@Z>`. While formally, these two are different type signatures,
looking at it from the programmers's perspective, each device in `Z` holds two
units of data `X` and `Y`. What happens here is not unlike `zip` in Python, and
indeed, we offer an operator `tff.federated_zip` that allows you to perform such
conversions explicity. When the `tff.federated_map` encounters a tuple as a
second argument, it simply invokes `tff.federated_zip` for you.
Given the above, you should now be able to recognize the expression
`tff.federated_broadcast(model)` as representing a value of TFF type
`{MODEL_TYPE}@CLIENTS`, and `data` as a value of TFF type
`{LOCAL_DATA_TYPE}@CLIENTS` (or simply `CLIENT_DATA_TYPE`), the two getting
filtered together through an implicit `tff.federated_zip` to form the second
argument to `tff.federated_map`.
The operator `tff.federated_broadcast`, as you'd expect, simply transfers data
from the server to the clients.
Now, let's see how our local training affected the average loss in the system.
%% Cell type:code id: tags:
```
print('initial_model loss =', federated_eval(initial_model,
federated_train_data))
print('locally_trained_model loss =',
federated_eval(locally_trained_model, federated_train_data))
```
%% Output
initial_model loss = 23.025852
locally_trained_model loss = 54.43263
%% Cell type:markdown id: tags:
Indeed, as expected, the loss has increased. In order to improve the model for
all users, we'll need to train in on everyone's data.
%% Cell type:markdown id: tags:
### Federated training
The simplest way to implement federated training is to locally train, and then
average the models. This uses the same building blocks and patters we've already
discussed, as you can see below.
%% Cell type:code id: tags:
```
SERVER_FLOAT_TYPE = tff.FederatedType(tf.float32, tff.SERVER)
@tff.federated_computation(SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE,
CLIENT_DATA_TYPE)
def federated_train(model, learning_rate, data):
return tff.federated_mean(
tff.federated_map(local_train, [
tff.federated_broadcast(model),
tff.federated_broadcast(learning_rate), data
]))
```
%% Cell type:markdown id: tags:
Note that in the full-featured implementation of Federated Averaging provided by
`tff.learning`, rather than averaging the models, we prefer to average model
deltas, for a number of reasons, e.g., the ability to clip the update norms,
for compression, etc.
Let's see whether the training works by running a few rounds of training and
comparing the average loss before and after.
%% Cell type:code id: tags:
```
model = initial_model
learning_rate = 0.1
for round_num in range(5):
model = federated_train(model, learning_rate, federated_train_data)
learning_rate = learning_rate * 0.9
loss = federated_eval(model, federated_train_data)
print('round {}, loss={}'.format(round_num, loss))
```
%% Output
round 0, loss=21.60552406311035
round 1, loss=20.365678787231445
round 2, loss=19.27480125427246
round 3, loss=18.31110954284668
round 4, loss=17.457256317138672
%% Cell type:markdown id: tags:
For completeness, let's now also run on the test data to confirm that our model
generalizes well.
%% Cell type:code id: tags:
```
print('initial_model test loss =',
federated_eval(initial_model, federated_test_data))
print('trained_model test loss =', federated_eval(model, federated_test_data))
```
%% Output
initial_model test loss = 22.795593
trained_model test loss = 17.278767
%% Cell type:markdown id: tags:
This concludes our tutorial.
Of course, our simplified example doesn't reflect a number of things you'd need
to do in a more realistic scenario - for example, we haven't computed metrics
other than loss. We encourage you to study
[the implementation](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/federated_averaging.py)
of federated averaging in `tff.learning` as a more complete example, and as a
way to demonstrate some of the coding practices we'd like to encourage.
......
docs/tutorials/federated_learning_for_image_classification.ipynb
浏览文件 @ 7c0121a5
%% Cell type:markdown id: tags:
 
##### Copyright 2019 The TensorFlow Authors.
 
%% Cell type:code id: tags:
 
```
#@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 obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
```
 
%% Cell type:markdown id: tags:
 
# Federated Learning for Image Classification
 
%% Cell type:markdown id: tags:
 
<table class="tfo-notebook-buttons" align="left">
<td>
<a target="_blank" href="https://www.tensorflow.org/federated/tutorials/federated_learning_for_image_classification"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
</td>
<td>
<a target="_blank" href="https://colab.research.google.com/github/tensorflow/federated/blob/v0.12.0/docs/tutorials/federated_learning_for_image_classification.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/federated_learning_for_image_classification.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
</td>
<td>
<a target="_blank" href="https://github.com/tensorflow/federated/blob/v0.12.0/docs/tutorials/federated_learning_for_image_classification.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/federated_learning_for_image_classification.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
</td>
</table>
 
%% Cell type:markdown id: tags:
 
**NOTE**: This colab has been verified to work with the [latest released version](https://github.com/tensorflow/federated#compatibility) of the `tensorflow_federated` pip package, but the Tensorflow Federated project is still in pre-release development and may not work on `master`.
 
In this tutorial, we use the classic MNIST training example to introduce the
Federated Learning (FL) API layer of TFF, `tff.learning` - a set of
higher-level interfaces that can be used to perform common types of federated
learning tasks, such as federated training, against user-supplied models
implemented in TensorFlow.
 
This tutorial, and the Federated Learning API, are intended primarily for users
who want to plug their own TensorFlow models into TFF, treating the latter
mostly as a black box. For a more in-depth understanding of TFF and how to
implement your own federated learning algorithms, see the tutorials on the FC Core API - [Custom Federated Algorithms Part 1](custom_federated_algorithms_1.ipynb) and [Part 2](custom_federated_algorithms_2.ipynb).
 
For more on `tff.learning`, continue with the
[Federated Learning for Text Generation](federated_learning_for_text_generation.ipynb),
tutorial which in addition to covering recurrent models, also demonstrates loading a
pre-trained serialized Keras model for refinement with federated learning
combined with evaluation using Keras.
 
%% Cell type:markdown id: tags:
 
## Before we start
 
Before we start, please run the following to make sure that your environment is
correctly setup. If you don't see a greeting, please refer to the
[Installation](../install.md) guide for instructions.
 
%% Cell type:code id: tags:
 
```
#@test {"skip": true}
!pip install --quiet --upgrade tensorflow_federated
 
# NOTE: Jupyter requires a patch to asyncio.
!pip install --quiet --upgrade nest_asyncio
import nest_asyncio
nest_asyncio.apply()
 
%load_ext tensorboard
```
 
%% Output
 
 |████████████████████████████████| 102kB 2.8MB/s
 |████████████████████████████████| 296kB 8.9MB/s
 |████████████████████████████████| 2.2MB 13.5MB/s
 |████████████████████████████████| 421.8MB 30kB/s
 |████████████████████████████████| 3.8MB 50.5MB/s
 |████████████████████████████████| 450kB 50.9MB/s
[?25h
 
%% Cell type:code id: tags:
 
```
import collections
 
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
 
tf.compat.v1.enable_v2_behavior()
 
np.random.seed(0)
 
tff.federated_computation(lambda: 'Hello, World!')()
```
 
%% Output
 
b'Hello, World!'
 
%% Cell type:markdown id: tags:
 
## Preparing the input data
 
Let's start with the data. Federated learning requires a federated data set,
i.e., a collection of data from multiple users. Federated data is typically
non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables),
which poses a unique set of challenges.
 
In order to facilitate experimentation, we seeded the TFF repository with a few
datasets, including a federated version of MNIST that contains a version of the [original NIST dataset](https://www.nist.gov/srd/nist-special-database-19) that has been re-processed using [Leaf](https://github.com/TalwalkarLab/leaf) so that the data is keyed by the original writer of the digits. Since each writer has a unique style, this dataset exhibits the kind of non-i.i.d. behavior expected of federated datasets.
 
Here's how we can load it.
 
%% Cell type:code id: tags:
 
```
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()
```
 
%% Output
 
Downloading data from https://storage.googleapis.com/tff-datasets-public/fed_emnist_digitsonly.tar.bz2
97402880/97398400 [==============================] - 1s 0us/step
 
%% Cell type:markdown id: tags:
 
The data sets returned by `load_data()` are instances of
`tff.simulation.ClientData`, an interface that allows you to enumerate the set
of users, to construct a `tf.data.Dataset` that represents the data of a
particular user, and to query the structure of individual elements. Here's how
you can use this interface to explore the content of the data set. Keep in mind
that while this interface allows you to iterate over clients ids, this is only a
feature of the simulation data. As you will see shortly, client identities are
not used by the federated learning framework - their only purpose is to allow
you to select subsets of the data for simulations.
 
%% Cell type:code id: tags:
 
```
len(emnist_train.client_ids)
```
 
%% Output
 
3383
 
%% Cell type:code id: tags:
 
```
emnist_train.element_type_structure
```
 
%% Output
 
OrderedDict([('label', TensorSpec(shape=(), dtype=tf.int32, name=None)),
('pixels',
TensorSpec(shape=(28, 28), dtype=tf.float32, name=None))])
 
%% Cell type:code id: tags:
 
```
example_dataset = emnist_train.create_tf_dataset_for_client(
emnist_train.client_ids[0])
 
example_element = next(iter(example_dataset))
 
example_element['label'].numpy()
```
 
%% Output
 
5
 
%% Cell type:code id: tags:
 
```
from matplotlib import pyplot as plt
plt.imshow(example_element['pixels'].numpy(), cmap='gray', aspect='equal')
plt.grid(False)
_ = plt.show()
```
 
%% Output
 
 
%% Cell type:markdown id: tags:
 
Since the data is already a `tf.data.Dataset`, preprocessing can be accomplished using Dataset transformations. Here, we flatten the `28x28` images
into `784`-element arrays, shuffle the individual examples, organize them into batches, and renames the features
from `pixels` and `label` to `x` and `y` for use with Keras. We also throw in a
`repeat` over the data set to run several epochs.
 
%% Cell type:code id: tags:
 
```
NUM_CLIENTS = 10
NUM_EPOCHS = 5
BATCH_SIZE = 20
SHUFFLE_BUFFER = 100
PREFETCH_BUFFER=10
 
def preprocess(dataset):
 
def batch_format_fn(element):
"""Flatten a batch `pixels` and return the features as an `OrderedDict`."""
return collections.OrderedDict(
x=tf.reshape(element['pixels'], [-1, 784]),
y=tf.reshape(element['label'], [-1, 1]))
 
return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(
BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)
```
 
%% Cell type:markdown id: tags:
 
Let's verify this worked.
 
%% Cell type:code id: tags:
 
```
preprocessed_example_dataset = preprocess(example_dataset)
 
sample_batch = tf.nest.map_structure(lambda x: x.numpy(),
next(iter(preprocessed_example_dataset)))
 
sample_batch
```
 
%% Output
 
OrderedDict([('x', array([[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]], dtype=float32)),
('y', array([[6],
[0],
[1],
[1],
[7],
[7],
[0],
[9],
[2],
[0],
[6],
[3],
[1],
[6],
[2],
[5],
[8],
[5],
[7],
[0]], dtype=int32))])
 
%% Cell type:markdown id: tags:
 
We have almost all the building blocks in place to construct federated data
sets.
 
One of the ways to feed federated data to TFF in a simulation is simply as a
Python list, with each element of the list holding the data of an individual
user, whether as a list or as a `tf.data.Dataset`. Since we already have
an interface that provides the latter, let's use it.
 
Here's a simple helper function that will construct a list of datasets from the
given set of users as an input to a round of training or evaluation.
 
%% Cell type:code id: tags:
 
```
def make_federated_data(client_data, client_ids):
return [
preprocess(client_data.create_tf_dataset_for_client(x))
for x in client_ids
]
```
 
%% Cell type:markdown id: tags:
 
Now, how do we choose clients?
 
In a typical federated training scenario, we are dealing with potentially a very
large population of user devices, only a fraction of which may be available for
training at a given point in time. This is the case, for example, when the
client devices are mobile phones that participate in training only when plugged
into a power source, off a metered network, and otherwise idle.
 
Of course, we are in a simulation environment, and all the data is locally
available. Typically then, when running simulations, we would simply sample a
random subset of the clients to be involved in each round of training, generally
different in each round.
 
That said, as you can find out by studying the paper on the
[Federated Averaging](https://arxiv.org/abs/1602.05629) algorithm, achieving convergence in a system with randomly sampled
subsets of clients in each round can take a while, and it would be impractical
to have to run hundreds of rounds in this interactive tutorial.
 
What we'll do instead is sample the set of clients once, and
reuse the same set across rounds to speed up convergence (intentionally
over-fitting to these few user's data). We leave it as an exercise for the
reader to modify this tutorial to simulate random sampling - it is fairly easy to
do (once you do, keep in mind that getting the model to converge may take a
while).
 
%% Cell type:code id: tags:
 
```
sample_clients = emnist_train.client_ids[0:NUM_CLIENTS]
 
federated_train_data = make_federated_data(emnist_train, sample_clients)
 
print('Number of client datasets: {l}'.format(l=len(federated_train_data)))
print('First dataset: {d}'.format(d=federated_train_data[0]))
```
 
%% Output
 
Number of client datasets: 10
First dataset: <PrefetchDataset shapes: OrderedDict([(x, (None, 784)), (y, (None, 1))]), types: OrderedDict([(x, tf.float32), (y, tf.int32)])>
 
%% Cell type:markdown id: tags:
 
## Creating a model with Keras
 
If you are using Keras, you likely already have code that constructs a Keras
model. Here's an example of a simple model that will suffice for our needs.
 
%% Cell type:code id: tags:
 
```
def create_keras_model():
return tf.keras.models.Sequential([
tf.keras.layers.Input(shape=(784,)),
tf.keras.layers.Dense(10, kernel_initializer='zeros'),
tf.keras.layers.Softmax(),
])
```
 
%% Cell type:markdown id: tags:
 
**Note:** we do not compile the model yet. The loss, metrics, and optimizers are introduced later.
 
In order to use any model with TFF, it needs to be wrapped in an instance of the
`tff.learning.Model` interface, which exposes methods to stamp the model's
forward pass, metadata properties, etc., similarly to Keras, but also introduces
additional elements, such as ways to control the process of computing federated
metrics. Let's not worry about this for now; if you have a Keras model like the
one we've just defined above, you can have TFF wrap it for you by invoking
`tff.learning.from_keras_model`, passing the model and a sample data batch as
arguments, as shown below.
 
%% Cell type:code id: tags:
 
```
def model_fn():
# We _must_ create a new model here, and _not_ capture it from an external
# scope. TFF will call this within different graph contexts.
keras_model = create_keras_model()
return tff.learning.from_keras_model(
keras_model,
dummy_batch=sample_batch,
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
```
 
%% Cell type:markdown id: tags:
 
## Training the model on federated data
 
Now that we have a model wrapped as `tff.learning.Model` for use with TFF, we
can let TFF construct a Federated Averaging algorithm by invoking the helper
function `tff.learning.build_federated_averaging_process`, as follows.
 
Keep in mind that the argument needs to be a constructor (such as `model_fn`
above), not an already-constructed instance, so that the construction of your
model can happen in a context controlled by TFF (if you're curious about the
reasons for this, we encourage you to read the follow-up tutorial on
[custom algorithms](custom_federated_algorithms_1.ipynb)).
 
One critical note on the Federated Averaging algorithm below, there are **2**
optimizers: a _client_optimizer_ and a _server_optimizer_. The
_client_optimizer_ is only used to compute local model updates on each client.
The _server_optimizer_ applies the averaged update to the global model at the
server. In particular, this means that the choice of optimizer and learning rate
used may need to be different than the ones you have used to train the model on
a standard i.i.d. dataset. We recommend starting with regular SGD, possibly with
a smaller learning rate than usual. The learning rate we use has not been
carefully tuned, feel free to experiment.
 
%% Cell type:code id: tags:
 
```
iterative_process = tff.learning.build_federated_averaging_process(
model_fn,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))
```
 
%% Cell type:markdown id: tags: