diff --git a/docs/federated_core.md b/docs/federated_core.md
index 79db0433f..a7851e6d3 100644
--- a/docs/federated_core.md
+++ b/docs/federated_core.md
@@ -87,17 +87,18 @@ though. Rather, we provide a Python API (the `tff` namespace) that wraps arounds
it as a way to define computations.
Specifically, TFF provides Python function decorators such as
-`tff.federated_computation` that trace the bodies of the decorated functions,
-and produce serialized representations of the federated computation logic in
-TFF's language. A function decorated with `tff.federated_computation` acts as a
-carrier of such serialized representation, and can embed it as a building block
-in the body of another computation, or execute it on demand when invoked.
+`federated_language.federated_computation` that trace the bodies of the
+decorated functions, and produce serialized representations of the federated
+computation logic in TFF's language. A function decorated with
+`federatedd_language.federated_computation` acts as a carrier of such serialized
+representation, and can embed it as a building block in the body of another
+computation, or execute it on demand when invoked.
Here's just one example; more examples can be found in the
[custom algorithms](tutorials/custom_federated_algorithms_1.ipynb) tutorials.
```python
-@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))
+@federated_language.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))
def get_average_temperature(sensor_readings):
return federated_language.federated_mean(sensor_readings)
```
@@ -285,18 +286,18 @@ public API:
For example, `federated_language.federated_broadcast` can be thought of as a
template operator of a functional type `T@SERVER -> T@CLIENTS`.
-* **Lambda expressions** (`tff.federated_computation`). A lambda expression in
- TFF is the equivalent of a `lambda` or `def` in Python; it consists of the
- parameter name, and a body (expression) that contains references to this
- parameter.
+* **Lambda expressions** (`federated_language.federated_computation`). A
+ lambda expression in TFF is the equivalent of a `lambda` or `def` in Python;
+ it consists of the parameter name, and a body (expression) that contains
+ references to this parameter.
In Python code, these can be created by decorating Python functions with
- `tff.federated_computation` and defining an argument.
+ `federated_language.federated_computation` and defining an argument.
Here's an example of a lambda expression we've already mentioned earlier:
```python
- @tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))
+ @federated_language.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))
def get_average_temperature(sensor_readings):
return federated_language.federated_mean(sensor_readings)
```
@@ -323,4 +324,4 @@ public API:
* Forming **tuples** and **selecting** their elements. Python expressions of
the form `[x, y]`, `x[y]`, or `x.y` that appear in the bodies of functions
- decorated with `tff.federated_computation`.
+ decorated with `federated_language.federated_computation`.
diff --git a/docs/install.md b/docs/install.md
index a55d13ceb..8e493a3c0 100644
--- a/docs/install.md
+++ b/docs/install.md
@@ -41,7 +41,7 @@ Note: To exit the virtual environment, run `deactivate`.
### 4. Test Tensorflow Federated.
-python -c "import tensorflow_federated as tff; print(tff.federated_computation(lambda: 'Hello World')())"
+python -c "import tensorflow_federated as tff; print(tff.tensorflow_computation(lambda: 'Hello World')())"
Success: The latest TensorFlow Federated Python package is now installed.
@@ -125,7 +125,7 @@ Note: To exit the virtual environment, run `deactivate`.
### 10. Test Tensorflow Federated.
-python -c "import tensorflow_federated as tff; print(tff.federated_computation(lambda: 'Hello World')())"
+python -c "import tensorflow_federated as tff; print(tff.tensorflow_computation(lambda: 'Hello World')())"
Success: A TensorFlow Federated Python package is now built from source and
diff --git a/docs/openmined2020/openmined_conference_2020.ipynb b/docs/openmined2020/openmined_conference_2020.ipynb
index 2a64e9db0..3c875e04d 100644
--- a/docs/openmined2020/openmined_conference_2020.ipynb
+++ b/docs/openmined2020/openmined_conference_2020.ipynb
@@ -31,13 +31,14 @@
},
"outputs": [],
"source": [
- "#@title Upgrade tensorflow_federated and load TensorBoard\n",
- "#@test {\"skip\": true}\n",
+ "# @title Upgrade tensorflow_federated and load TensorBoard\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated\n",
"!pip install --quiet --upgrade nest-asyncio\n",
"\n",
"import nest_asyncio\n",
+ "\n",
"nest_asyncio.apply()\n",
"\n",
"%load_ext tensorboard\n",
@@ -45,8 +46,9 @@
"import sys\n",
"\n",
"if not sys.warnoptions:\n",
- " import warnings\n",
- " warnings.simplefilter(\"ignore\")"
+ " import warnings\n",
+ "\n",
+ " warnings.simplefilter(\"ignore\")"
]
},
{
@@ -57,12 +59,11 @@
},
"outputs": [],
"source": [
- "#@title\n",
+ "# @title\n",
"import collections\n",
- "from matplotlib import pyplot as plt\n",
- "from IPython.display import display, HTML, IFrame\n",
- "\n",
"import federated_language\n",
+ "from IPython.display import HTML, IFrame, display\n",
+ "from matplotlib import pyplot as plt\n",
"import numpy as np\n",
"import tensorflow as tf\n",
"import tensorflow_federated as tff\n",
@@ -71,10 +72,18 @@
"\n",
"np.random.seed(0)\n",
"\n",
+ "\n",
"def greetings():\n",
- " display(HTML('\u003cb\u003e\u003cfont size=\"6\" color=\"#ff00f4\"\u003eGreetings, virtual tutorial participants!\u003c/font\u003e\u003c/b\u003e'))\n",
+ " display(\n",
+ " HTML(\n",
+ " '\u003cb\u003e\u003cfont size=\"6\" color=\"#ff00f4\"\u003eGreetings, virtual tutorial'\n",
+ " ' participants!\u003c/font\u003e\u003c/b\u003e'\n",
+ " )\n",
+ " )\n",
" return True\n",
- "l = tff.federated_computation(greetings)()"
+ "\n",
+ "\n",
+ "l = federated_language.federated_computation(greetings)()"
]
},
{
@@ -205,7 +214,8 @@
"source": [
"# Let's look at the shape of our data\n",
"example_dataset = emnist_train.create_tf_dataset_for_client(\n",
- " emnist_train.client_ids[0])\n",
+ " emnist_train.client_ids[0]\n",
+ ")\n",
"\n",
"example_dataset.element_spec"
]
@@ -220,7 +230,8 @@
"source": [
"# Let's select an example dataset from one of our simulated clients\n",
"example_dataset = emnist_train.create_tf_dataset_for_client(\n",
- " emnist_train.client_ids[0])\n",
+ " emnist_train.client_ids[0]\n",
+ ")\n",
"\n",
"# Your code to get an example element from one client:\n",
"example_element = next(iter(example_dataset))\n",
@@ -259,11 +270,11 @@
"outputs": [],
"source": [
"## Example MNIST digits for one client\n",
- "f = plt.figure(figsize=(20,4))\n",
+ "f = plt.figure(figsize=(20, 4))\n",
"j = 0\n",
"\n",
"for e in example_dataset.take(40):\n",
- " plt.subplot(4, 10, j+1)\n",
+ " plt.subplot(4, 10, j + 1)\n",
" plt.imshow(e['pixels'].numpy(), cmap='gray', aspect='equal')\n",
" plt.axis('off')\n",
" j += 1"
@@ -278,17 +289,17 @@
"outputs": [],
"source": [
"# Number of examples per layer for a sample of clients\n",
- "f = plt.figure(figsize=(12,7))\n",
+ "f = plt.figure(figsize=(12, 7))\n",
"f.suptitle(\"Label Counts for a Sample of Clients\")\n",
"for i in range(6):\n",
" ds = emnist_train.create_tf_dataset_for_client(emnist_train.client_ids[i])\n",
" k = collections.defaultdict(list)\n",
" for e in ds:\n",
- " k[e['label'].numpy()].append(e['label'].numpy())\n",
- " plt.subplot(2, 3, i+1)\n",
+ " k[e[\"label\"].numpy()].append(e[\"label\"].numpy())\n",
+ " plt.subplot(2, 3, i + 1)\n",
" plt.title(\"Client {}\".format(i))\n",
" for j in range(10):\n",
- " plt.hist(k[j], density=False, bins=[0,1,2,3,4,5,6,7,8,9,10])\n"
+ " plt.hist(k[j], density=False, bins=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])"
]
},
{
@@ -307,12 +318,12 @@
" k = collections.defaultdict(list)\n",
" for e in ds:\n",
" k[e['label'].numpy()].append(e['pixels'].numpy())\n",
- " f = plt.figure(i, figsize=(12,5))\n",
+ " f = plt.figure(i, figsize=(12, 5))\n",
" f.suptitle(\"Client #{}'s Mean Image Per Label\".format(i))\n",
" for j in range(10):\n",
- " mn_img = np.mean(k[j],0)\n",
- " plt.subplot(2, 5, j+1)\n",
- " plt.imshow(mn_img.reshape((28,28)))#,cmap='gray')\n",
+ " mn_img = np.mean(k[j], 0)\n",
+ " plt.subplot(2, 5, j + 1)\n",
+ " plt.imshow(mn_img.reshape((28, 28))) # ,cmap='gray')\n",
" plt.axis('off')\n",
"\n",
"# Each client has different mean images -- each client will be nudging the model\n",
@@ -342,7 +353,8 @@
"NUM_EPOCHS = 5\n",
"BATCH_SIZE = 20\n",
"SHUFFLE_BUFFER = 100\n",
- "PREFETCH_BUFFER=10\n",
+ "PREFETCH_BUFFER = 10\n",
+ "\n",
"\n",
"def preprocess(dataset):\n",
"\n",
@@ -350,10 +362,16 @@
" \"\"\"Flatten a batch `pixels` and return the features as an `OrderedDict`.\"\"\"\n",
" return collections.OrderedDict(\n",
" x=tf.reshape(element['pixels'], [-1, 784]),\n",
- " y=tf.reshape(element['label'], [-1, 1]))\n",
+ " y=tf.reshape(element['label'], [-1, 1]),\n",
+ " )\n",
"\n",
- " return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(\n",
- " BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)"
+ " return (\n",
+ " dataset.repeat(NUM_EPOCHS)\n",
+ " .shuffle(SHUFFLE_BUFFER)\n",
+ " .batch(BATCH_SIZE)\n",
+ " .map(batch_format_fn)\n",
+ " .prefetch(PREFETCH_BUFFER)\n",
+ " )"
]
},
{
@@ -375,8 +393,9 @@
"source": [
"preprocessed_example_dataset = preprocess(example_dataset)\n",
"\n",
- "sample_batch = tf.nest.map_structure(lambda x: x.numpy(),\n",
- " next(iter(preprocessed_example_dataset)))\n",
+ "sample_batch = tf.nest.map_structure(\n",
+ " lambda x: x.numpy(), next(iter(preprocessed_example_dataset))\n",
+ ")\n",
"\n",
"sample_batch"
]
@@ -492,14 +511,9 @@
"mod.compile(\n",
" optimizer=tf.keras.optimizers.RMSprop(),\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
- " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]\n",
- ")\n",
- "h = mod.fit(\n",
- " x_train,\n",
- " y_train,\n",
- " batch_size=64,\n",
- " epochs=2\n",
+ " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],\n",
")\n",
+ "h = mod.fit(x_train, y_train, batch_size=64, epochs=2)\n",
"\n",
"# ------------------------------------------------------------------------------"
]
@@ -548,7 +562,8 @@
" keras_model,\n",
" input_spec=preprocessed_example_dataset.element_spec,\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
- " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
+ " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],\n",
+ " )"
]
},
{
@@ -576,8 +591,8 @@
" model_fn,\n",
" client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),\n",
" # Add server optimizer here!\n",
- " server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))\n",
- ""
+ " server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0),\n",
+ ")"
]
},
{
@@ -713,7 +728,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"import os\n",
"import shutil\n",
"\n",
@@ -744,7 +759,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"with summary_writer.as_default():\n",
" for round_num in range(1, NUM_ROUNDS):\n",
" state, metrics = iterative_process.next(state, federated_train_data)\n",
@@ -769,7 +784,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"%tensorboard --logdir /tmp/logs/scalars/ --port=0"
]
},
@@ -827,6 +842,7 @@
"outputs": [],
"source": [
"import random\n",
+ "\n",
"shuffled_ids = emnist_test.client_ids.copy()\n",
"random.shuffle(shuffled_ids)\n",
"sample_clients = shuffled_ids[0:NUM_CLIENTS]\n",
@@ -923,12 +939,15 @@
"NUM_CLIENTS = 10\n",
"BATCH_SIZE = 20\n",
"\n",
+ "\n",
"def preprocess(dataset):\n",
"\n",
" def batch_format_fn(element):\n",
" \"\"\"Flatten a batch of EMNIST data and return a (features, label) tuple.\"\"\"\n",
- " return (tf.reshape(element['pixels'], [-1, 784]),\n",
- " tf.reshape(element['label'], [-1, 1]))\n",
+ " return (\n",
+ " tf.reshape(element['pixels'], [-1, 784]),\n",
+ " tf.reshape(element['label'], [-1, 1]),\n",
+ " )\n",
"\n",
" return dataset.batch(BATCH_SIZE).map(batch_format_fn)"
]
@@ -941,10 +960,12 @@
},
"outputs": [],
"source": [
- "client_ids = np.random.choice(emnist_train.client_ids, size=NUM_CLIENTS, replace=False)\n",
+ "client_ids = np.random.choice(\n",
+ " emnist_train.client_ids, size=NUM_CLIENTS, replace=False\n",
+ ")\n",
"\n",
- "federated_train_data = [preprocess(emnist_train.create_tf_dataset_for_client(x))\n",
- " for x in client_ids\n",
+ "federated_train_data = [\n",
+ " preprocess(emnist_train.create_tf_dataset_for_client(x)) for x in client_ids\n",
"]"
]
},
@@ -1005,7 +1026,8 @@
" keras_model,\n",
" input_spec=federated_train_data[0].element_spec,\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
- " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
+ " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],\n",
+ " )"
]
},
{
@@ -1145,8 +1167,9 @@
" # Initialize the client model with the current server weights.\n",
" client_weights = model.weights.trainable\n",
" # Assign the server weights to the client model.\n",
- " tf.nest.map_structure(lambda x, y: x.assign(y),\n",
- " client_weights, server_weights)\n",
+ " tf.nest.map_structure(\n",
+ " lambda x, y: x.assign(y), client_weights, server_weights\n",
+ " )\n",
"\n",
" # Use the client_optimizer to update the local model.\n",
" for batch in dataset:\n",
@@ -1188,8 +1211,9 @@
" \"\"\"Updates the server model weights as the average of the client model weights.\"\"\"\n",
" model_weights = model.weights.trainable\n",
" # Assign the mean client weights to the server model.\n",
- " tf.nest.map_structure(lambda x, y: x.assign(y),\n",
- " model_weights, mean_client_weights)\n",
+ " tf.nest.map_structure(\n",
+ " lambda x, y: x.assign(y), model_weights, mean_client_weights\n",
+ " )\n",
" return model_weights"
]
},
@@ -1249,7 +1273,9 @@
},
"outputs": [],
"source": [
- "federated_float_on_clients = federated_language.FederatedType(np.float32, federated_language.CLIENTS)"
+ "federated_float_on_clients = federated_language.FederatedType(\n",
+ " np.float32, federated_language.CLIENTS\n",
+ ")"
]
},
{
@@ -1311,7 +1337,9 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(np.float32, federated_language.CLIENTS)\n",
+ ")\n",
"def get_average_temperature(client_temperatures):\n",
" return federated_language.federated_mean(client_temperatures)"
]
@@ -1322,7 +1350,7 @@
"id": "iSgs6Te5Lo8r"
},
"source": [
- "You might ask, how is this different from the `tf.function` decorator in TensorFlow? The key answer is that the code generated by `tff.federated_computation` is neither TensorFlow nor Python code; It is a specification of a distributed system in an internal platform-independent *glue language*.\n",
+ "You might ask, how is this different from the `tf.function` decorator in TensorFlow? The key answer is that the code generated by `federated_language.federated_computation` is neither TensorFlow nor Python code; It is a specification of a distributed system in an internal platform-independent *glue language*.\n",
"\n",
"While this may sound complicated, you can think of TFF computations as functions with well-defined type signatures. These type signatures can be directly queried."
]
@@ -1344,9 +1372,9 @@
"id": "TveOYFfuLo8s"
},
"source": [
- "This `tff.federated_computation` accepts arguments of federated type `{float32}@CLIENTS`, and returns values of federated type `{float32}@SERVER`. Federated computations may also go from server to client, from client to client, or from server to server. Federated computations can also be composed like normal functions, as long as their type signatures match up.\n",
+ "This `federated_language.federated_computation` accepts arguments of federated type `{float32}@CLIENTS`, and returns values of federated type `{float32}@SERVER`. Federated computations may also go from server to client, from client to client, or from server to server. Federated computations can also be composed like normal functions, as long as their type signatures match up.\n",
"\n",
- "To support development, TFF allows you to invoke a `tff.federated_computation` as a Python function. For example, we can call"
+ "To support development, TFF allows you to invoke a `federated_language.federated_computation` as a Python function. For example, we can call"
]
},
{
@@ -1375,7 +1403,7 @@
"id": "nwyj8f3HLo8w"
},
"source": [
- "There are two key restrictions to be aware of. First, when the Python interpreter encounters a `tff.federated_computation` decorator, the function is traced once and serialized for future use. Therefore, TFF computations are fundamentally *non-eager*. This behavior is somewhat analogous to that of the [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function) decorator in TensorFlow.\n",
+ "There are two key restrictions to be aware of. First, when the Python interpreter encounters a `federated_language.federated_computation` decorator, the function is traced once and serialized for future use. Therefore, TFF computations are fundamentally *non-eager*. This behavior is somewhat analogous to that of the [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function) decorator in TensorFlow.\n",
"\n",
"Second, a federated computation can only consist of federated operators (such as `federated_language.federated_mean`), they cannot contain TensorFlow operations. TensorFlow code must be confined to blocks decorated with `tff.tensorflow.computation`. Most ordinary TensorFlow code can be directly decotrated, such as the following function that takes a number and adds `0.5` to it."
]
@@ -1419,7 +1447,7 @@
"id": "WNjwrNMjLo8z"
},
"source": [
- "Here we see an important difference between `tff.federated_computation` and `tff.tensorflow.computation`. The former has explicit placements, while the latter does not.\n",
+ "Here we see an important difference between `federated_language.federated_computation` and `tff.tensorflow.computation`. The former has explicit placements, while the latter does not.\n",
"\n",
"We can use `tff.tensorflow.computation` blocks in federated computations by specifying placements. Let's create a function that adds half, but only to federated floats at the clients. We can do this by using `federated_language.federated_map`, which applies a given `tff.tensorflow.computation`, while preserving the placement."
]
@@ -1432,7 +1460,9 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(np.float32, federated_language.CLIENTS)\n",
+ ")\n",
"def add_half_on_clients(x):\n",
" return federated_language.federated_map(add_half, x)"
]
@@ -1467,7 +1497,7 @@
"\n",
"* TFF operates on federated values.\n",
"* Each federated value has a *federated type*, with a *type* (eg. `np.float32`) and a *placement* (eg. `federated_language.CLIENTS`).\n",
- "* Federated values can be transformed using *federated computations*, which must be decorated with `tff.federated_computation` and a federated type signature.\n",
+ "* Federated values can be transformed using *federated computations*, which must be decorated with `federated_language.federated_computation` and a federated type signature.\n",
"* TensorFlow code must be contained in blocks with `tff.tensorflow.computation` decorators.\n",
"* These blocks can then be incorporated into federated computations.\n"
]
@@ -1482,7 +1512,7 @@
"\n",
"Now that we've peeked at the Federated Core, we can build our own federated learning algorithm. Remember that above, we defined an `initialize_fn` and `next_fn` for our algorithm. The `next_fn` will make use of the `client_update` and `server_update` we defined using pure TensorFlow code.\n",
"\n",
- "However, in order to make our algorithm a federated computation, we will need both the `next_fn` and `initialize_fn` to be `tff.federated_computations`."
+ "However, in order to make our algorithm a federated computation, we will need both the `next_fn` and `initialize_fn` to be `federated_language.federated_computation`s."
]
},
{
@@ -1536,9 +1566,11 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation\n",
+ "@federated_language.federated_computation\n",
"def initialize_fn():\n",
- " return federated_language.federated_value(server_init(), federated_language.SERVER)"
+ " return federated_language.federated_value(\n",
+ " server_init(), federated_language.SERVER\n",
+ " )"
]
},
{
@@ -1679,7 +1711,7 @@
"id": "SImhLbu4Lo9D"
},
"source": [
- "Last, but not least, we need to create the `tff.federated_computation` that brings this all together. This function will accept two *federated values*, one corresponding to the server weights (with placement `federated_language.SERVER`), and the other corresponding to the client datasets (with placement `federated_language.CLIENTS`).\n",
+ "Last, but not least, we need to create the `federated_language.federated_computation` that brings this all together. This function will accept two *federated values*, one corresponding to the server weights (with placement `federated_language.SERVER`), and the other corresponding to the client datasets (with placement `federated_language.CLIENTS`).\n",
"\n",
"Note that both these types were defined above! We simply need to give them the proper placement using `federated_language.type_at_{server/clients}``."
]
@@ -1692,8 +1724,12 @@
},
"outputs": [],
"source": [
- "federated_server_type = federated_language.FederatedType(model_weights_type, federated_language.SERVER)\n",
- "federated_dataset_type = federated_language.FederatedType(tf_dataset_type, federated_language.CLIENTS)"
+ "federated_server_type = federated_language.FederatedType(\n",
+ " model_weights_type, federated_language.SERVER\n",
+ ")\n",
+ "federated_dataset_type = federated_language.FederatedType(\n",
+ " tf_dataset_type, federated_language.CLIENTS\n",
+ ")"
]
},
{
@@ -1720,20 +1756,27 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_server_type, federated_dataset_type)\n",
+ "@federated_language.federated_computation(\n",
+ " federated_server_type, federated_dataset_type\n",
+ ")\n",
"def next_fn(server_weights, federated_dataset):\n",
" # Broadcast the server weights to the clients.\n",
- " server_weights_at_client = federated_language.federated_broadcast(server_weights)\n",
+ " server_weights_at_client = federated_language.federated_broadcast(\n",
+ " server_weights\n",
+ " )\n",
"\n",
" # Each client computes their updated weights.\n",
" client_weights = federated_language.federated_map(\n",
- " client_update_fn, (federated_dataset, server_weights_at_client))\n",
+ " client_update_fn, (federated_dataset, server_weights_at_client)\n",
+ " )\n",
"\n",
" # The server averages these updates.\n",
" mean_client_weights = federated_language.federated_mean(client_weights)\n",
"\n",
" # The server updates its model.\n",
- " server_weights = federated_language.federated_map(server_update_fn, mean_client_weights)\n",
+ " server_weights = federated_language.federated_map(\n",
+ " server_update_fn, mean_client_weights\n",
+ " )\n",
"\n",
" return server_weights"
]
@@ -1744,7 +1787,7 @@
"id": "kWomG3TtLo9I"
},
"source": [
- "We now have a `tff.federated_computation` for both the algorithm initialization, and for running one step of the algorithm. To finish our algorithm, we pass these into `tff.templates.IterativeProcess`."
+ "We now have a `federated_language.federated_computation` for both the algorithm initialization, and for running one step of the algorithm. To finish our algorithm, we pass these into `tff.templates.IterativeProcess`."
]
},
{
@@ -1756,8 +1799,7 @@
"outputs": [],
"source": [
"federated_algorithm = tff.templates.IterativeProcess(\n",
- " initialize_fn=initialize_fn,\n",
- " next_fn=next_fn\n",
+ " initialize_fn=initialize_fn, next_fn=next_fn\n",
")"
]
},
@@ -1840,7 +1882,9 @@
},
"outputs": [],
"source": [
- "central_emnist_test = emnist_test.create_tf_dataset_from_all_clients().take(1000)\n",
+ "central_emnist_test = emnist_test.create_tf_dataset_from_all_clients().take(\n",
+ " 1000\n",
+ ")\n",
"central_emnist_test = preprocess(central_emnist_test)"
]
},
@@ -1865,7 +1909,7 @@
" keras_model = create_keras_model()\n",
" keras_model.compile(\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
- " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]\n",
+ " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],\n",
" )\n",
" keras_model.set_weights(server_state)\n",
" keras_model.evaluate(central_emnist_test)"
diff --git a/docs/tff_for_research.md b/docs/tff_for_research.md
index 56ab08350..1b3708fea 100644
--- a/docs/tff_for_research.md
+++ b/docs/tff_for_research.md
@@ -28,8 +28,8 @@ types of logic.
individual `tf.function`s from 1. by wrapping them as
`tff.tensorflow.computation`s and then orchestrating them using abstractions
like `federated_language.federated_broadcast` and
- `federated_language.federated_mean` inside a `tff.federated_computation`.
- See, for example, this
+ `federated_language.federated_mean` inside a
+ `federated_language.federated_computation`. See, for example, this
[orchestration for Federated Averaging](https://github.com/google-parfait/tensorflow-federated/blob/main/examples/simple_fedavg/simple_fedavg_tff.py#L112-L140).
1. An outer driver script that simulates the control logic of a production FL
diff --git a/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb b/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb
index cd87611e0..4bfd2ffd2 100644
--- a/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb
+++ b/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -74,7 +74,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -169,12 +169,15 @@
"NUM_CLIENTS = 10\n",
"BATCH_SIZE = 20\n",
"\n",
+ "\n",
"def preprocess(dataset):\n",
"\n",
" def batch_format_fn(element):\n",
" \"\"\"Flatten a batch of EMNIST data and return a (features, label) tuple.\"\"\"\n",
- " return (tf.reshape(element['pixels'], [-1, 784]),\n",
- " tf.reshape(element['label'], [-1, 1]))\n",
+ " return (\n",
+ " tf.reshape(element['pixels'], [-1, 784]),\n",
+ " tf.reshape(element['label'], [-1, 1]),\n",
+ " )\n",
"\n",
" return dataset.batch(BATCH_SIZE).map(batch_format_fn)"
]
@@ -197,8 +200,8 @@
"outputs": [],
"source": [
"client_ids = sorted(emnist_train.client_ids)[:NUM_CLIENTS]\n",
- "federated_train_data = [preprocess(emnist_train.create_tf_dataset_for_client(x))\n",
- " for x in client_ids\n",
+ "federated_train_data = [\n",
+ " preprocess(emnist_train.create_tf_dataset_for_client(x)) for x in client_ids\n",
"]"
]
},
@@ -324,7 +327,7 @@
"outputs": [],
"source": [
"def initialize_fn():\n",
- " trainable_weights, _ = tff_model.initial_weights\n",
+ " trainable_weights, _ = tff_model.initial_weights\n",
" return trainable_weights"
]
},
@@ -516,7 +519,9 @@
},
"outputs": [],
"source": [
- "federated_float_on_clients = federated_language.FederatedType(np.float32, federated_language.CLIENTS)"
+ "federated_float_on_clients = federated_language.FederatedType(\n",
+ " np.float32, federated_language.CLIENTS\n",
+ ")"
]
},
{
@@ -593,7 +598,9 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(np.float32, federated_language.CLIENTS)\n",
+ ")\n",
"def get_average_temperature(client_temperatures):\n",
" return federated_language.federated_mean(client_temperatures)"
]
@@ -604,7 +611,7 @@
"id": "iSgs6Te5Lo8r"
},
"source": [
- "You might ask, how is this different from the `tf.function` decorator in TensorFlow? The key answer is that the code generated by `tff.federated_computation` is neither TensorFlow nor Python code; It is a specification of a distributed system in an internal platform-independent *glue language*.\n",
+ "You might ask, how is this different from the `tf.function` decorator in TensorFlow? The key answer is that the code generated by `federated_language.federated_computation` is neither TensorFlow nor Python code; It is a specification of a distributed system in an internal platform-independent *glue language*.\n",
"\n",
"While this may sound complicated, you can think of TFF computations as functions with well-defined type signatures. These type signatures can be directly queried."
]
@@ -641,9 +648,9 @@
"id": "TveOYFfuLo8s"
},
"source": [
- "This `tff.federated_computation` accepts arguments of federated type `{float32}@CLIENTS`, and returns values of federated type `{float32}@SERVER`. Federated computations may also go from server to client, from client to client, or from server to server. Federated computations can also be composed like normal functions, as long as their type signatures match up.\n",
+ "This `federated_language.federated_computation` accepts arguments of federated type `{float32}@CLIENTS`, and returns values of federated type `{float32}@SERVER`. Federated computations may also go from server to client, from client to client, or from server to server. Federated computations can also be composed like normal functions, as long as their type signatures match up.\n",
"\n",
- "To support development, TFF allows you to invoke a `tff.federated_computation` as a Python function. For example, you can call"
+ "To support development, TFF allows you to invoke a `federated_language.federated_computation` as a Python function. For example, you can call"
]
},
{
@@ -684,7 +691,7 @@
"id": "nwyj8f3HLo8w"
},
"source": [
- "There are two key restrictions to be aware of. First, when the Python interpreter encounters a `tff.federated_computation` decorator, the function is traced once and serialized for future use. Due to the decentralized nature of Federated Learning, this future usage may occur elsewhere, such as a remote execution environment. Therefore, TFF computations are fundamentally *non-eager*. This behavior is somewhat analogous to that of the [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function) decorator in TensorFlow.\n",
+ "There are two key restrictions to be aware of. First, when the Python interpreter encounters a `federated_language.federated_computation` decorator, the function is traced once and serialized for future use. Due to the decentralized nature of Federated Learning, this future usage may occur elsewhere, such as a remote execution environment. Therefore, TFF computations are fundamentally *non-eager*. This behavior is somewhat analogous to that of the [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function) decorator in TensorFlow.\n",
"\n",
"Second, a federated computation can only consist of federated operators (such as `federated_language.federated_mean`), they cannot contain TensorFlow operations. TensorFlow code must be confined to blocks decorated with `tff.tensorflow.computation`. Most ordinary TensorFlow code can be directly decorated, such as the following function that takes a number and adds `0.5` to it."
]
@@ -743,7 +750,7 @@
"id": "WNjwrNMjLo8z"
},
"source": [
- "This showcases an important difference between `tff.federated_computation` and `tff.tensorflow.computation`. The former has explicit placements, while the latter does not.\n",
+ "This showcases an important difference between `federated_language.federated_computation` and `tff.tensorflow.computation`. The former has explicit placements, while the latter does not.\n",
"\n",
"You can use `tff.tensorflow.computation` blocks in federated computations by specifying placements. Let's create a function that adds half, but only to federated floats at the clients. You can do this by using `federated_language.federated_map`, which applies a given `tff.tensorflow.computation`, while preserving the placement."
]
@@ -756,7 +763,9 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(np.float32, federated_language.CLIENTS)\n",
+ ")\n",
"def add_half_on_clients(x):\n",
" return federated_language.federated_map(add_half, x)"
]
@@ -806,7 +815,7 @@
"\n",
"* TFF operates on federated values.\n",
"* Each federated value has a *federated type*, with a *type* (eg. `np.float32`) and a *placement* (eg. `federated_language.CLIENTS`).\n",
- "* Federated values can be transformed using *federated computations*, which must be decorated with `tff.federated_computation` and a federated type signature.\n",
+ "* Federated values can be transformed using *federated computations*, which must be decorated with `federated_language.federated_computation` and a federated type signature.\n",
"* TensorFlow code must be contained in blocks with `tff.tensorflow.computation` decorators.\n",
"* These blocks can then be incorporated into federated computations.\n"
]
@@ -821,7 +830,7 @@
"\n",
"Now that you've gotten a glimpse of the Federated Core, you can build our own federated learning algorithm. Remember that above, you defined an `initialize_fn` and `next_fn` for our algorithm. The `next_fn` will make use of the `client_update` and `server_update` you defined using pure TensorFlow code.\n",
"\n",
- "However, in order to make our algorithm a federated computation, you will need both the `next_fn` and `initialize_fn` to each be a `tff.federated_computation`."
+ "However, in order to make our algorithm a federated computation, you will need both the `next_fn` and `initialize_fn` to each be a `federated_language.federated_computation`."
]
},
{
@@ -874,9 +883,11 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation\n",
+ "@federated_language.federated_computation\n",
"def initialize_fn():\n",
- " return federated_language.federated_eval(server_init, federated_language.SERVER)"
+ " return federated_language.federated_eval(\n",
+ " server_init, federated_language.SERVER\n",
+ " )"
]
},
{
@@ -1046,7 +1057,7 @@
"id": "SImhLbu4Lo9D"
},
"source": [
- "Last, but not least, you need to create the `tff.federated_computation` that brings this all together. This function will accept two *federated values*, one corresponding to the server weights (with placement `federated_language.SERVER`), and the other corresponding to the client datasets (with placement `federated_language.CLIENTS`).\n",
+ "Last, but not least, you need to create the `federated_language.federated_computation` that brings this all together. This function will accept two *federated values*, one corresponding to the server weights (with placement `federated_language.SERVER`), and the other corresponding to the client datasets (with placement `federated_language.CLIENTS`).\n",
"\n",
"Note that both these types were defined above! You simply need to give them the proper placement using `federated_language.FederatedType`."
]
@@ -1059,8 +1070,12 @@
},
"outputs": [],
"source": [
- "federated_server_type = federated_language.FederatedType(model_weights_type, federated_language.SERVER)\n",
- "federated_dataset_type = federated_language.FederatedType(tf_dataset_type, federated_language.CLIENTS)"
+ "federated_server_type = federated_language.FederatedType(\n",
+ " model_weights_type, federated_language.SERVER\n",
+ ")\n",
+ "federated_dataset_type = federated_language.FederatedType(\n",
+ " tf_dataset_type, federated_language.CLIENTS\n",
+ ")"
]
},
{
@@ -1087,10 +1102,14 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_server_type, federated_dataset_type)\n",
+ "@federated_language.federated_computation(\n",
+ " federated_server_type, federated_dataset_type\n",
+ ")\n",
"def next_fn(server_weights, federated_dataset):\n",
" # Broadcast the server weights to the clients.\n",
- " server_weights_at_client = federated_language.federated_broadcast(server_weights)\n",
+ " server_weights_at_client = federated_language.federated_broadcast(\n",
+ " server_weights\n",
+ " )\n",
"\n",
" # Each client computes their updated weights.\n",
" client_weights = federated_language.federated_map(\n",
@@ -1101,7 +1120,9 @@
" mean_client_weights = federated_language.federated_mean(client_weights)\n",
"\n",
" # The server updates its model.\n",
- " server_weights = federated_language.federated_map(server_update_fn, mean_client_weights)\n",
+ " server_weights = federated_language.federated_map(\n",
+ " server_update_fn, mean_client_weights\n",
+ " )\n",
"\n",
" return server_weights"
]
@@ -1112,7 +1133,7 @@
"id": "kWomG3TtLo9I"
},
"source": [
- "You now have a `tff.federated_computation` for both the algorithm initialization, and for running one step of the algorithm. To finish our algorithm, you pass these into `tff.templates.IterativeProcess`."
+ "You now have a `federated_language.federated_computation` for both the algorithm initialization, and for running one step of the algorithm. To finish our algorithm, you pass these into `tff.templates.IterativeProcess`."
]
},
{
@@ -1124,8 +1145,7 @@
"outputs": [],
"source": [
"federated_algorithm = tff.templates.IterativeProcess(\n",
- " initialize_fn=initialize_fn,\n",
- " next_fn=next_fn\n",
+ " initialize_fn=initialize_fn, next_fn=next_fn\n",
")"
]
},
diff --git a/docs/tutorials/composing_learning_algorithms.ipynb b/docs/tutorials/composing_learning_algorithms.ipynb
index 2a15c6369..09643b6c7 100644
--- a/docs/tutorials/composing_learning_algorithms.ipynb
+++ b/docs/tutorials/composing_learning_algorithms.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -83,7 +83,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -291,15 +291,19 @@
" optimizer: tff.learning.optimizers.Optimizer,\n",
") -\u003e tff.learning.templates.ClientWorkProcess:\n",
" \"\"\"Creates a client work process that uses gradient clipping.\"\"\"\n",
- " data_type = federated_language.SequenceType(tff.tensorflow.to_type(model.input_spec))\n",
+ " data_type = federated_language.SequenceType(\n",
+ " tff.tensorflow.to_type(model.input_spec)\n",
+ " )\n",
" model_weights_type = federated_language.to_type(\n",
" tf.nest.map_structure(\n",
- " lambda arr: federated_language.TensorType(shape=arr.shape, dtype=arr.dtype),\n",
+ " lambda arr: federated_language.TensorType(\n",
+ " shape=arr.shape, dtype=arr.dtype\n",
+ " ),\n",
" tff.learning.models.ModelWeights(*model.initial_weights),\n",
" )\n",
" )\n",
"\n",
- " @tff.federated_computation\n",
+ " @federated_language.federated_computation\n",
" def initialize_fn():\n",
" return federated_language.federated_value((), federated_language.SERVER)\n",
"\n",
@@ -307,9 +311,11 @@
" def client_update_computation(model_weights, dataset):\n",
" return client_update(model, dataset, model_weights, optimizer)\n",
"\n",
- " @tff.federated_computation(\n",
+ " @federated_language.federated_computation(\n",
" initialize_fn.type_signature.result,\n",
- " federated_language.FederatedType(model_weights_type, federated_language.CLIENTS),\n",
+ " federated_language.FederatedType(\n",
+ " model_weights_type, federated_language.CLIENTS\n",
+ " ),\n",
" federated_language.FederatedType(data_type, federated_language.CLIENTS),\n",
" )\n",
" def next_fn(state, model_weights, client_dataset):\n",
@@ -318,7 +324,9 @@
" )\n",
" # Return empty measurements, though a more complete algorithm might\n",
" # measure something here.\n",
- " measurements = federated_language.federated_value((), federated_language.SERVER)\n",
+ " measurements = federated_language.federated_value(\n",
+ " (), federated_language.SERVER\n",
+ " )\n",
" return tff.templates.MeasuredProcessOutput(\n",
" state, client_result, measurements\n",
" )\n",
@@ -380,18 +388,22 @@
"NUM_CLIENTS = 10\n",
"BATCH_SIZE = 20\n",
"\n",
+ "\n",
"def preprocess(dataset):\n",
"\n",
" def batch_format_fn(element):\n",
" \"\"\"Flatten a batch of EMNIST data and return a (features, label) tuple.\"\"\"\n",
- " return (tf.reshape(element['pixels'], [-1, 784]),\n",
- " tf.reshape(element['label'], [-1, 1]))\n",
+ " return (\n",
+ " tf.reshape(element['pixels'], [-1, 784]),\n",
+ " tf.reshape(element['label'], [-1, 1]),\n",
+ " )\n",
"\n",
" return dataset.batch(BATCH_SIZE).map(batch_format_fn)\n",
"\n",
+ "\n",
"client_ids = sorted(emnist_train.client_ids)[:NUM_CLIENTS]\n",
- "federated_train_data = [preprocess(emnist_train.create_tf_dataset_for_client(x))\n",
- " for x in client_ids\n",
+ "federated_train_data = [\n",
+ " preprocess(emnist_train.create_tf_dataset_for_client(x)) for x in client_ids\n",
"]"
]
},
@@ -529,11 +541,7 @@
"outputs": [],
"source": [
"fed_avg_with_clipping = tff.learning.templates.compose_learning_process(\n",
- " initial_model_weights_fn,\n",
- " distributor,\n",
- " client_work,\n",
- " aggregator,\n",
- " finalizer\n",
+ " initial_model_weights_fn, distributor, client_work, aggregator, finalizer\n",
")"
]
},
@@ -622,7 +630,9 @@
},
"outputs": [],
"source": [
- "learning_process_output = fed_avg_with_clipping.next(state, federated_train_data)"
+ "learning_process_output = fed_avg_with_clipping.next(\n",
+ " state, federated_train_data\n",
+ ")"
]
},
{
diff --git a/docs/tutorials/custom_aggregators.ipynb b/docs/tutorials/custom_aggregators.ipynb
index e8cd41f8a..60c7bfe86 100644
--- a/docs/tutorials/custom_aggregators.ipynb
+++ b/docs/tutorials/custom_aggregators.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -70,7 +70,7 @@
"source": [
"In this tutorial, we explain design principles behind the `tff.aggregators` module and best practices for implementing custom aggregation of values from clients to server.\n",
"\n",
- "**Prerequisites.** This tutorial assumes you are already familiar with basic concepts of [Federated Core](https://www.tensorflow.org/federated/federated_core) such as placements (`federated_language.SERVER`, `federated_language.CLIENTS`), how TFF represents computations (`tff.tensorflow.computation`, `tff.federated_computation`) and their type signatures."
+ "**Prerequisites.** This tutorial assumes you are already familiar with basic concepts of [Federated Core](https://www.tensorflow.org/federated/federated_core) such as placements (`federated_language.SERVER`, `federated_language.CLIENTS`), how TFF represents computations (`tff.tensorflow.computation`, `federated_language.federated_computation`) and their type signatures."
]
},
{
@@ -81,7 +81,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -341,21 +341,30 @@
"class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
"\n",
" def create(self, value_type):\n",
- " @tff.federated_computation()\n",
+ " @federated_language.federated_computation()\n",
" def initialize_fn():\n",
" return federated_language.federated_value((), federated_language.SERVER)\n",
"\n",
- " @tff.federated_computation(initialize_fn.type_signature.result,\n",
- " federated_language.FederatedType(value_type, federated_language.CLIENTS))\n",
+ " @federated_language.federated_computation(\n",
+ " initialize_fn.type_signature.result,\n",
+ " federated_language.FederatedType(\n",
+ " value_type, federated_language.CLIENTS\n",
+ " ),\n",
+ " )\n",
" def next_fn(state, value):\n",
" scaled_value = federated_language.federated_map(\n",
- " tff.tensorflow.computation(lambda x: x * 2.0), value)\n",
+ " tff.tensorflow.computation(lambda x: x * 2.0), value\n",
+ " )\n",
" summed_value = federated_language.federated_sum(scaled_value)\n",
" unscaled_value = federated_language.federated_map(\n",
- " tff.tensorflow.computation(lambda x: x / 2.0), summed_value)\n",
- " measurements = federated_language.federated_value((), federated_language.SERVER)\n",
+ " tff.tensorflow.computation(lambda x: x / 2.0), summed_value\n",
+ " )\n",
+ " measurements = federated_language.federated_value(\n",
+ " (), federated_language.SERVER\n",
+ " )\n",
" return tff.templates.MeasuredProcessOutput(\n",
- " state=state, result=unscaled_value, measurements=measurements)\n",
+ " state=state, result=unscaled_value, measurements=measurements\n",
+ " )\n",
"\n",
" return tff.templates.AggregationProcess(initialize_fn, next_fn)"
]
@@ -393,9 +402,11 @@
"client_data = [1.0, 2.0, 5.0]\n",
"factory = ExampleTaskFactory()\n",
"aggregation_process = factory.create(federated_language.TensorType(np.float32))\n",
- "print(f'Type signatures of the created aggregation process:\\n'\n",
- " f' - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
- " f' - next: {aggregation_process.next.type_signature}\\n')\n",
+ "print(\n",
+ " 'Type signatures of the created aggregation process:\\n'\n",
+ " f' - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
+ " f' - next: {aggregation_process.next.type_signature}\\n'\n",
+ ")\n",
"\n",
"state = aggregation_process.initialize()\n",
"output = aggregation_process.next(state, client_data)\n",
@@ -448,23 +459,33 @@
"class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
"\n",
" def create(self, value_type):\n",
- " @tff.federated_computation()\n",
+ " @federated_language.federated_computation()\n",
" def initialize_fn():\n",
" return federated_language.federated_value(0.0, federated_language.SERVER)\n",
"\n",
- " @tff.federated_computation(initialize_fn.type_signature.result,\n",
- " federated_language.FederatedType(value_type, federated_language.CLIENTS))\n",
+ " @federated_language.federated_computation(\n",
+ " initialize_fn.type_signature.result,\n",
+ " federated_language.FederatedType(\n",
+ " value_type, federated_language.CLIENTS\n",
+ " ),\n",
+ " )\n",
" def next_fn(state, value):\n",
" new_state = federated_language.federated_map(\n",
- " tff.tensorflow.computation(lambda x: x + 1.0), state)\n",
+ " tff.tensorflow.computation(lambda x: x + 1.0), state\n",
+ " )\n",
" state_at_clients = federated_language.federated_broadcast(new_state)\n",
" scaled_value = federated_language.federated_map(\n",
- " tff.tensorflow.computation(lambda x, y: x * y), (value, state_at_clients))\n",
+ " tff.tensorflow.computation(lambda x, y: x * y),\n",
+ " (value, state_at_clients),\n",
+ " )\n",
" summed_value = federated_language.federated_sum(scaled_value)\n",
" unscaled_value = federated_language.federated_map(\n",
- " tff.tensorflow.computation(lambda x, y: x / y), (summed_value, new_state))\n",
+ " tff.tensorflow.computation(lambda x, y: x / y),\n",
+ " (summed_value, new_state),\n",
+ " )\n",
" return tff.templates.MeasuredProcessOutput(\n",
- " state=new_state, result=unscaled_value, measurements=summed_value)\n",
+ " state=new_state, result=unscaled_value, measurements=summed_value\n",
+ " )\n",
"\n",
" return tff.templates.AggregationProcess(initialize_fn, next_fn)"
]
@@ -514,9 +535,11 @@
"client_data = [1.0, 2.0, 5.0]\n",
"factory = ExampleTaskFactory()\n",
"aggregation_process = factory.create(federated_language.TensorType(np.float32))\n",
- "print(f'Type signatures of the created aggregation process:\\n'\n",
- " f' - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
- " f' - next: {aggregation_process.next.type_signature}\\n')\n",
+ "print(\n",
+ " 'Type signatures of the created aggregation process:\\n'\n",
+ " f' - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
+ " f' - next: {aggregation_process.next.type_signature}\\n'\n",
+ ")\n",
"\n",
"state = aggregation_process.initialize()\n",
"\n",
@@ -572,31 +595,43 @@
"def scale(value, factor):\n",
" return tf.nest.map_structure(lambda x: x * factor, value)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation()\n",
"def unscale(value, factor):\n",
" return tf.nest.map_structure(lambda x: x / factor, value)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation()\n",
"def add_one(value):\n",
" return value + 1.0\n",
"\n",
+ "\n",
"class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
"\n",
" def create(self, value_type):\n",
- " @tff.federated_computation()\n",
+ " @federated_language.federated_computation()\n",
" def initialize_fn():\n",
" return federated_language.federated_value(0.0, federated_language.SERVER)\n",
"\n",
- " @tff.federated_computation(initialize_fn.type_signature.result,\n",
- " federated_language.FederatedType(value_type, federated_language.CLIENTS))\n",
+ " @federated_language.federated_computation(\n",
+ " initialize_fn.type_signature.result,\n",
+ " federated_language.FederatedType(\n",
+ " value_type, federated_language.CLIENTS\n",
+ " ),\n",
+ " )\n",
" def next_fn(state, value):\n",
" new_state = federated_language.federated_map(add_one, state)\n",
" state_at_clients = federated_language.federated_broadcast(new_state)\n",
- " scaled_value = federated_language.federated_map(scale, (value, state_at_clients))\n",
+ " scaled_value = federated_language.federated_map(\n",
+ " scale, (value, state_at_clients)\n",
+ " )\n",
" summed_value = federated_language.federated_sum(scaled_value)\n",
- " unscaled_value = federated_language.federated_map(unscale, (summed_value, new_state))\n",
+ " unscaled_value = federated_language.federated_map(\n",
+ " unscale, (summed_value, new_state)\n",
+ " )\n",
" return tff.templates.MeasuredProcessOutput(\n",
- " state=new_state, result=unscaled_value, measurements=summed_value)\n",
+ " state=new_state, result=unscaled_value, measurements=summed_value\n",
+ " )\n",
"\n",
" return tff.templates.AggregationProcess(initialize_fn, next_fn)"
]
@@ -607,7 +642,7 @@
"id": "x5-1gxv1uebz"
},
"source": [
- "This example highlights a pattern which may be useful to follow when structuring TFF code. When not dealing with very simple operations, the code becomes more legible when the `tff.tensorflow.computation`s that will be used as building blocks inside a `tff.federated_computation` are created in a separate place. Inside of the `tff.federated_computation`, these building blocks are only connected using the intrinsic operators."
+ "This example highlights a pattern which may be useful to follow when structuring TFF code. When not dealing with very simple operations, the code becomes more legible when the `tff.tensorflow.computation`s that will be used as building blocks inside a `federated_language.federated_computation` are created in a separate place. Inside of the `federated_language.federated_computation`, these building blocks are only connected using the intrinsic operators."
]
},
{
@@ -641,19 +676,23 @@
}
],
"source": [
- "client_data = [[[1.0, 2.0], [3.0, 4.0, 5.0]],\n",
- " [[1.0, 1.0], [3.0, 0.0, -5.0]]]\n",
+ "client_data = [[[1.0, 2.0], [3.0, 4.0, 5.0]], [[1.0, 1.0], [3.0, 0.0, -5.0]]]\n",
"factory = ExampleTaskFactory()\n",
"aggregation_process = factory.create(\n",
- " federated_language.to_type([(np.float32, (2,)), (np.float32, (3,))]))\n",
- "print(f'Type signatures of the created aggregation process:\\n'\n",
- " f' - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
- " f' - next: {aggregation_process.next.type_signature}\\n')\n",
+ " federated_language.to_type([(np.float32, (2,)), (np.float32, (3,))])\n",
+ ")\n",
+ "print(\n",
+ " 'Type signatures of the created aggregation process:\\n'\n",
+ " f' - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
+ " f' - next: {aggregation_process.next.type_signature}\\n'\n",
+ ")\n",
"\n",
"state = aggregation_process.initialize()\n",
"output = aggregation_process.next(state, client_data)\n",
- "print(f'Aggregation result: [{output.result[0]}, {output.result[1]}]\\n'\n",
- " f' Expected: [[2. 3.], [6. 4. 0.]]')"
+ "print(\n",
+ " f'Aggregation result: [{output.result[0]}, {output.result[1]}]\\n'\n",
+ " ' Expected: [[2. 3.], [6. 4. 0.]]'\n",
+ ")"
]
},
{
@@ -696,14 +735,17 @@
"def scale(value, factor):\n",
" return tf.nest.map_structure(lambda x: x * factor, value)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation()\n",
"def unscale(value, factor):\n",
" return tf.nest.map_structure(lambda x: x / factor, value)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation()\n",
"def add_one(value):\n",
" return value + 1.0\n",
"\n",
+ "\n",
"class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
"\n",
" def __init__(self, inner_factory=None):\n",
@@ -714,33 +756,48 @@
" def create(self, value_type):\n",
" inner_process = self._inner_factory.create(value_type)\n",
"\n",
- " @tff.federated_computation()\n",
+ " @federated_language.federated_computation()\n",
" def initialize_fn():\n",
- " my_state = federated_language.federated_value(0.0, federated_language.SERVER)\n",
+ " my_state = federated_language.federated_value(\n",
+ " 0.0, federated_language.SERVER\n",
+ " )\n",
" inner_state = inner_process.initialize()\n",
" return federated_language.federated_zip((my_state, inner_state))\n",
"\n",
- " @tff.federated_computation(initialize_fn.type_signature.result,\n",
- " federated_language.FederatedType(value_type, federated_language.CLIENTS))\n",
+ " @federated_language.federated_computation(\n",
+ " initialize_fn.type_signature.result,\n",
+ " federated_language.FederatedType(\n",
+ " value_type, federated_language.CLIENTS\n",
+ " ),\n",
+ " )\n",
" def next_fn(state, value):\n",
" my_state, inner_state = state\n",
" my_new_state = federated_language.federated_map(add_one, my_state)\n",
" my_state_at_clients = federated_language.federated_broadcast(my_new_state)\n",
- " scaled_value = federated_language.federated_map(scale, (value, my_state_at_clients))\n",
+ " scaled_value = federated_language.federated_map(\n",
+ " scale, (value, my_state_at_clients)\n",
+ " )\n",
"\n",
" # Delegation to an inner factory, returning values placed at SERVER.\n",
" inner_output = inner_process.next(inner_state, scaled_value)\n",
"\n",
- " unscaled_value = federated_language.federated_map(unscale, (inner_output.result, my_new_state))\n",
+ " unscaled_value = federated_language.federated_map(\n",
+ " unscale, (inner_output.result, my_new_state)\n",
+ " )\n",
"\n",
- " new_state = federated_language.federated_zip((my_new_state, inner_output.state))\n",
+ " new_state = federated_language.federated_zip(\n",
+ " (my_new_state, inner_output.state)\n",
+ " )\n",
" measurements = federated_language.federated_zip(\n",
" collections.OrderedDict(\n",
" scaled_value=inner_output.result,\n",
- " example_task=inner_output.measurements))\n",
+ " example_task=inner_output.measurements,\n",
+ " )\n",
+ " )\n",
"\n",
" return tff.templates.MeasuredProcessOutput(\n",
- " state=new_state, result=unscaled_value, measurements=measurements)\n",
+ " state=new_state, result=unscaled_value, measurements=measurements\n",
+ " )\n",
"\n",
" return tff.templates.AggregationProcess(initialize_fn, next_fn)"
]
@@ -800,14 +857,22 @@
"output = aggregation_process.next(state, client_data)\n",
"print('| Round #1')\n",
"print(f'| Aggregation result: {output.result} (expected 8.0)')\n",
- "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
- "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')\n",
+ "print(\n",
+ " f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}'\n",
+ ")\n",
+ "print(\n",
+ " f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}'\n",
+ ")\n",
"\n",
"output = aggregation_process.next(output.state, client_data)\n",
"print('\\n| Round #2')\n",
"print(f'| Aggregation result: {output.result} (expected 8.0)')\n",
- "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
- "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')"
+ "print(\n",
+ " f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}'\n",
+ ")\n",
+ "print(\n",
+ " f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}'\n",
+ ")"
]
},
{
@@ -855,14 +920,22 @@
"output = aggregation_process.next(state, client_data)\n",
"print('| Round #1')\n",
"print(f'| Aggregation result: {output.result} (expected 8.0)')\n",
- "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
- "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')\n",
+ "print(\n",
+ " f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}'\n",
+ ")\n",
+ "print(\n",
+ " f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}'\n",
+ ")\n",
"\n",
"output = aggregation_process.next(output.state, client_data)\n",
"print('\\n| Round #2')\n",
"print(f'| Aggregation result: {output.result} (expected 8.0)')\n",
- "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
- "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')"
+ "print(\n",
+ " f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}'\n",
+ ")\n",
+ "print(\n",
+ " f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}'\n",
+ ")"
]
},
{
diff --git a/docs/tutorials/custom_federated_algorithms_1.ipynb b/docs/tutorials/custom_federated_algorithms_1.ipynb
index dacbbcff4..bb5a7fcab 100644
--- a/docs/tutorials/custom_federated_algorithms_1.ipynb
+++ b/docs/tutorials/custom_federated_algorithms_1.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -170,7 +170,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -213,10 +213,11 @@
}
],
"source": [
- "@tff.federated_computation\n",
+ "@federated_language.federated_computation\n",
"def hello_world():\n",
" return 'Hello, World!'\n",
"\n",
+ "\n",
"hello_world()"
]
},
@@ -260,7 +261,9 @@
},
"outputs": [],
"source": [
- "federated_float_on_clients = federated_language.FederatedType(np.float32, federated_language.CLIENTS)"
+ "federated_float_on_clients = federated_language.FederatedType(\n",
+ " np.float32, federated_language.CLIENTS\n",
+ ")"
]
},
{
@@ -447,7 +450,11 @@
}
],
"source": [
- "str(federated_language.FederatedType(np.float32, federated_language.CLIENTS, all_equal=True))"
+ "str(\n",
+ " federated_language.FederatedType(\n",
+ " np.float32, federated_language.CLIENTS, all_equal=True\n",
+ " )\n",
+ ")"
]
},
{
@@ -494,8 +501,9 @@
}
],
"source": [
- "simple_regression_model_type = (\n",
- " federated_language.StructType([('a', np.float32), ('b', np.float32)]))\n",
+ "simple_regression_model_type = federated_language.StructType(\n",
+ " [('a', np.float32), ('b', np.float32)]\n",
+ ")\n",
"\n",
"str(simple_regression_model_type)"
]
@@ -536,8 +544,11 @@
}
],
"source": [
- "str(federated_language.FederatedType(\n",
- " simple_regression_model_type, federated_language.CLIENTS, all_equal=True))"
+ "str(\n",
+ " federated_language.FederatedType(\n",
+ " simple_regression_model_type, federated_language.CLIENTS, all_equal=True\n",
+ " )\n",
+ ")"
]
},
{
@@ -699,7 +710,9 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(np.float32, federated_language.CLIENTS)\n",
+ ")\n",
"def get_average_temperature(sensor_readings):\n",
" return federated_language.federated_mean(sensor_readings)"
]
@@ -716,7 +729,7 @@
"in TensorFlow, and if so, why introduce yet another one, and how is it\n",
"different?\n",
"\n",
- "The short answer is that the code generated by the `tff.federated_computation`\n",
+ "The short answer is that the code generated by the `federated_language.federated_computation`\n",
"wrapper is *neither* TensorFlow, *nor is it* Python - it's a specification of a\n",
"distributed system in an internal platform-independent *glue* language. At this\n",
"point, this will undoubtedly sound cryptic, but please bear this intuitive\n",
@@ -846,15 +859,15 @@
"you supplied client values at input, and consumed the server result.\n",
"\n",
"Now, let's return to a note we made earlier about the\n",
- "`tff.federated_computation` decorator emitting code in a *glue* language.\n",
+ "`federated_language.federated_computation` decorator emitting code in a *glue* language.\n",
"Although the logic of TFF computations can be expressed as ordinary functions in\n",
- "Python (you just need to decorate them with `tff.federated_computation` as we've\n",
+ "Python (you just need to decorate them with `federated_language.federated_computation` as we've\n",
"done above), and you can directly invoke them with Python arguments just\n",
"like any other Python functions in this notebook, behind the scenes, as we noted\n",
"earlier, TFF computations are actually *not* Python.\n",
"\n",
"What we mean by this is that when the Python interpreter encounters a function\n",
- "decorated with `tff.federated_computation`, it traces the statements in this\n",
+ "decorated with `federated_language.federated_computation`, it traces the statements in this\n",
"function's body once (at definition time), and then constructs a\n",
"[serialized representation](https://github.com/tensorflow/federated/blob/main/tensorflow_federated/proto/v0/computation.proto)\n",
"of the computation's logic for future use - whether for execution, or to be\n",
@@ -880,11 +893,16 @@
}
],
"source": [
- "@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(np.float32, federated_language.CLIENTS)\n",
+ ")\n",
"def get_average_temperature(sensor_readings):\n",
"\n",
- " print ('Getting traced, the argument is \"{}\".'.format(\n",
- " type(sensor_readings).__name__))\n",
+ " print(\n",
+ " 'Getting traced, the argument is \"{}\".'.format(\n",
+ " type(sensor_readings).__name__\n",
+ " )\n",
+ " )\n",
"\n",
" return federated_language.federated_mean(sensor_readings)"
]
@@ -913,7 +931,7 @@
"be aware of its existence, the fact that TFF computations are fundamentally\n",
"non-eager, and cannot capture arbitrary Python state. Python code contained in a\n",
"TFF computation's body is executed at definition time, when the body of the\n",
- "Python function decorated with `tff.federated_computation` is traced before\n",
+ "Python function decorated with `federated_language.federated_computation` is traced before\n",
"getting serialized. It's not retraced again at invocation time (except when the\n",
"function is polymorphic; please refer to the documentation pages for details).\n",
"\n",
@@ -950,29 +968,29 @@
"\n",
"The same is true of composing federated computations. The computation\n",
"`get_average_temperature` may be invoked in a body of another Python function\n",
- "decorated with `tff.federated_computation` - doing so will cause it to be\n",
+ "decorated with `federated_language.federated_computation` - doing so will cause it to be\n",
"embedded in the body of the parent, much in the same way `federated_language.federated_mean`\n",
"was embedded in its own body earlier.\n",
"\n",
"An important restriction to be aware of is that bodies of Python functions\n",
- "decorated with `tff.federated_computation` must consist *only* of federated\n",
+ "decorated with `federated_language.federated_computation` must consist *only* of federated\n",
"operators, i.e., they cannot directly contain TensorFlow operations. For\n",
"example, you cannot directly use `tf.nest` interfaces to add a pair of\n",
"federated values. TensorFlow code must be confined to blocks of code decorated\n",
"with a `tff.tensorflow.computation` discussed in the following section. Only when\n",
"wrapped in this manner can the wrapped TensorFlow code be invoked in the body of\n",
- "a `tff.federated_computation`.\n",
+ "a `federated_language.federated_computation`.\n",
"\n",
"The reasons for this separation are technical (it's hard to trick operators such\n",
"as `tf.add` to work with non-tensors) as well as architectural. The language of\n",
"federated computations (i.e., the logic constructed from serialized bodies of\n",
- "Python functions decorated with `tff.federated_computation`) is designed to\n",
+ "Python functions decorated with `federated_language.federated_computation`) is designed to\n",
"serve as a platform-independent *glue* language. This glue language is currently\n",
"used to build distributed systems from embedded sections of TensorFlow code\n",
"(confined to `tff.tensorflow.computation` blocks). In the fullness of time, we\n",
"anticipate the need to embed sections of other, non-TensorFlow logic, such as\n",
"relational database queries that might represent input pipelines, all connected\n",
- "together using the same glue language (the `tff.federated_computation` blocks)."
+ "together using the same glue language (the `federated_language.federated_computation` blocks)."
]
},
{
@@ -1084,7 +1102,9 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(federated_language.FederatedType(np.float32, federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(np.float32, federated_language.CLIENTS)\n",
+ ")\n",
"def add_half_on_clients(x):\n",
" return federated_language.federated_map(add_half, x)"
]
@@ -1123,7 +1143,7 @@
"### Executing TensorFlow computations\n",
"\n",
"Execution of computations defined with `tff.tensorflow.computation` follows the same\n",
- "rules as those we described for `tff.federated_computation`. They can be invoked\n",
+ "rules as those we described for `federated_language.federated_computation`. They can be invoked\n",
"as ordinary callables in Python, as follows."
]
},
@@ -1175,7 +1195,7 @@
"serialized representation.\n",
"\n",
"The only difference between Python methods decorated with\n",
- "`tff.federated_computation` and those decorated with `tff.tensorflow.computation` is\n",
+ "`federated_language.federated_computation` and those decorated with `tff.tensorflow.computation` is\n",
"that the latter are serialized as TensorFlow graphs (whereas the former are not\n",
"allowed to contain TensorFlow code directly embedded in them).\n",
"\n",
@@ -1209,14 +1229,14 @@
"try:\n",
"\n",
" # Eager mode\n",
- " constant_10 = tf.constant(10.)\n",
+ " constant_10 = tf.constant(10.0)\n",
"\n",
" @tff.tensorflow.computation(np.float32)\n",
" def add_ten(x):\n",
" return x + constant_10\n",
"\n",
"except Exception as err:\n",
- " print (err)"
+ " print(err)"
]
},
{
@@ -1256,12 +1276,14 @@
],
"source": [
"def get_constant_10():\n",
- " return tf.constant(10.)\n",
+ " return tf.constant(10.0)\n",
+ "\n",
"\n",
"@tff.tensorflow.computation(np.float32)\n",
"def add_ten(x):\n",
" return x + get_constant_10()\n",
"\n",
+ "\n",
"add_ten(5.0)"
]
},
@@ -1336,8 +1358,9 @@
"source": [
"@tff.tensorflow.computation(federated_language.SequenceType(np.float32))\n",
"def get_local_temperature_average(local_temperatures):\n",
- " sum_and_count = (\n",
- " local_temperatures.reduce((0.0, 0), lambda x, y: (x[0] + y, x[1] + 1)))\n",
+ " sum_and_count = local_temperatures.reduce(\n",
+ " (0.0, 0), lambda x, y: (x[0] + y, x[1] + 1)\n",
+ " )\n",
" return sum_and_count[0] / tf.cast(sum_and_count[1], tf.float32)"
]
},
@@ -1407,6 +1430,7 @@
"def foo(x):\n",
" return x.reduce(np.int32(0), lambda x, y: x + y)\n",
"\n",
+ "\n",
"foo([1, 2, 3])"
]
},
@@ -1486,7 +1510,11 @@
}
],
"source": [
- "@tff.tensorflow.computation(federated_language.SequenceType(collections.OrderedDict([('A', np.int32), ('B', np.int32)])))\n",
+ "@tff.tensorflow.computation(\n",
+ " federated_language.SequenceType(\n",
+ " collections.OrderedDict([('A', np.int32), ('B', np.int32)])\n",
+ " )\n",
+ ")\n",
"def foo(ds):\n",
" print('element_structure = {}'.format(ds.element_spec))\n",
" return ds.reduce(np.int32(0), lambda total, x: total + x['A'] * x['B'])"
@@ -1568,11 +1596,17 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(\n",
- " federated_language.FederatedType(federated_language.SequenceType(np.float32), federated_language.CLIENTS))\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(\n",
+ " federated_language.SequenceType(np.float32), federated_language.CLIENTS\n",
+ " )\n",
+ ")\n",
"def get_global_temperature_average(sensor_readings):\n",
" return federated_language.federated_mean(\n",
- " federated_language.federated_map(get_local_temperature_average, sensor_readings))"
+ " federated_language.federated_map(\n",
+ " get_local_temperature_average, sensor_readings\n",
+ " )\n",
+ " )"
]
},
{
diff --git a/docs/tutorials/custom_federated_algorithms_2.ipynb b/docs/tutorials/custom_federated_algorithms_2.ipynb
index 86aa04a63..3c83059da 100644
--- a/docs/tutorials/custom_federated_algorithms_2.ipynb
+++ b/docs/tutorials/custom_federated_algorithms_2.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -108,7 +108,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -149,10 +149,11 @@
}
],
"source": [
- "@tff.federated_computation\n",
+ "@federated_language.federated_computation\n",
"def hello_world():\n",
" return 'Hello, World!'\n",
"\n",
+ "\n",
"hello_world()"
]
},
@@ -256,13 +257,13 @@
" output_sequence = []\n",
" all_samples = [i for i, d in enumerate(source[1]) if d == digit]\n",
" for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):\n",
- " batch_samples = all_samples[i:i + BATCH_SIZE]\n",
+ " batch_samples = all_samples[i : i + BATCH_SIZE]\n",
" output_sequence.append({\n",
- " 'x':\n",
- " np.array([source[0][i].flatten() / 255.0 for i in batch_samples],\n",
- " dtype=np.float32),\n",
- " 'y':\n",
- " np.array([source[1][i] for i in batch_samples], dtype=np.int32)\n",
+ " 'x': np.array(\n",
+ " [source[0][i].flatten() / 255.0 for i in batch_samples],\n",
+ " dtype=np.float32,\n",
+ " ),\n",
+ " 'y': np.array([source[1][i] for i in batch_samples], dtype=np.int32),\n",
" })\n",
" return output_sequence\n",
"\n",
@@ -405,7 +406,8 @@
"source": [
"BATCH_SPEC = collections.OrderedDict(\n",
" x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),\n",
- " y=tf.TensorSpec(shape=[None], dtype=tf.int32))\n",
+ " y=tf.TensorSpec(shape=[None], dtype=tf.int32),\n",
+ ")\n",
"BATCH_TYPE = tff.tensorflow.to_type(BATCH_SPEC)\n",
"\n",
"str(BATCH_TYPE)"
@@ -454,7 +456,8 @@
"source": [
"MODEL_SPEC = collections.OrderedDict(\n",
" weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),\n",
- " bias=tf.TensorSpec(shape=[10], dtype=tf.float32))\n",
+ " bias=tf.TensorSpec(shape=[10], dtype=tf.float32),\n",
+ ")\n",
"MODEL_TYPE = tff.tensorflow.to_type(MODEL_SPEC)\n",
"print(MODEL_TYPE)"
]
@@ -480,13 +483,18 @@
"# be later called from within another tf.function. Necessary because a\n",
"# @tf.function decorated method cannot invoke a @tff.tensorflow.computation.\n",
"\n",
+ "\n",
"@tf.function\n",
"def forward_pass(model, batch):\n",
" predicted_y = tf.nn.softmax(\n",
- " tf.matmul(batch['x'], model['weights']) + model['bias'])\n",
+ " tf.matmul(batch['x'], model['weights']) + model['bias']\n",
+ " )\n",
" return -tf.reduce_mean(\n",
" tf.reduce_sum(\n",
- " tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))\n",
+ " tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]\n",
+ " )\n",
+ " )\n",
+ "\n",
"\n",
"@tff.tensorflow.computation(MODEL_TYPE, BATCH_TYPE)\n",
"def batch_loss(model, batch):\n",
@@ -563,7 +571,8 @@
"source": [
"initial_model = collections.OrderedDict(\n",
" weights=np.zeros([784, 10], dtype=np.float32),\n",
- " bias=np.zeros([10], dtype=np.float32))\n",
+ " bias=np.zeros([10], dtype=np.float32),\n",
+ ")\n",
"\n",
"sample_batch = federated_train_data[5][-1]\n",
"\n",
@@ -637,7 +646,8 @@
" loss = forward_pass(model_vars, batch)\n",
" grads = tape.gradient(loss, model_vars)\n",
" optimizer.apply_gradients(\n",
- " zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))\n",
+ " zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars))\n",
+ " )\n",
" return model_vars\n",
"\n",
" return _train_on_batch(model_vars, batch)"
@@ -751,11 +761,16 @@
"source": [
"LOCAL_DATA_TYPE = federated_language.SequenceType(BATCH_TYPE)\n",
"\n",
- "@tff.federated_computation(MODEL_TYPE, np.float32, LOCAL_DATA_TYPE)\n",
+ "\n",
+ "@federated_language.federated_computation(\n",
+ " MODEL_TYPE, np.float32, LOCAL_DATA_TYPE\n",
+ ")\n",
"def local_train(initial_model, learning_rate, all_batches):\n",
"\n",
" # Reduction function to apply to each batch.\n",
- " @tff.federated_computation((MODEL_TYPE, np.float32), BATCH_TYPE)\n",
+ " @federated_language.federated_computation(\n",
+ " (MODEL_TYPE, np.float32), BATCH_TYPE\n",
+ " )\n",
" def batch_fn(model_with_lr, batch):\n",
" model, lr = model_with_lr\n",
" return batch_train(model, batch, lr), lr\n",
@@ -804,7 +819,7 @@
"First, while we could have implemented this logic entirely in TensorFlow,\n",
"relying on `tf.data.Dataset.reduce` to process the sequence similarly to how\n",
"we've done it earlier, we've opted this time to express the logic in the glue\n",
- "language, as a `tff.federated_computation`. We've used the federated operator\n",
+ "language, as a `federated_language.federated_computation`. We've used the federated operator\n",
"`federated_language.sequence_reduce` to perform the reduction.\n",
"\n",
"The operator `federated_language.sequence_reduce` is used similarly to\n",
@@ -877,7 +892,7 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)\n",
+ "@federated_language.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)\n",
"def local_eval(model, all_batches):\n",
"\n",
" @tff.tensorflow.computation((MODEL_TYPE, np.float32), BATCH_TYPE)\n",
@@ -944,7 +959,7 @@
"Second, note that just as in `local_train`, the component function we need\n",
"(`batch_loss`) takes more parameters than what the federated operator\n",
"(`federated_language.sequence_map`) expects, so we again define a partial, this time inline by\n",
- "directly wrapping a `lambda` as a `tff.federated_computation`. Using wrappers\n",
+ "directly wrapping a `lambda` as a `federated_language.federated_computation`. Using wrappers\n",
"inline with a function as an argument is the recommended way to use\n",
"`tff.tensorflow.computation` to embed TensorFlow logic in TFF.\n",
"\n",
@@ -969,10 +984,13 @@
}
],
"source": [
- "print('initial_model loss =', local_eval(initial_model,\n",
- " federated_train_data[5]))\n",
- "print('locally_trained_model loss =',\n",
- " local_eval(locally_trained_model, federated_train_data[5]))"
+ "print(\n",
+ " 'initial_model loss =', local_eval(initial_model, federated_train_data[5])\n",
+ ")\n",
+ "print(\n",
+ " 'locally_trained_model loss =',\n",
+ " local_eval(locally_trained_model, federated_train_data[5]),\n",
+ ")"
]
},
{
@@ -1003,10 +1021,13 @@
}
],
"source": [
- "print('initial_model loss =', local_eval(initial_model,\n",
- " federated_train_data[0]))\n",
- "print('locally_trained_model loss =',\n",
- " local_eval(locally_trained_model, federated_train_data[0]))"
+ "print(\n",
+ " 'initial_model loss =', local_eval(initial_model, federated_train_data[0])\n",
+ ")\n",
+ "print(\n",
+ " 'locally_trained_model loss =',\n",
+ " local_eval(locally_trained_model, federated_train_data[0]),\n",
+ ")"
]
},
{
@@ -1042,8 +1063,12 @@
},
"outputs": [],
"source": [
- "SERVER_MODEL_TYPE = federated_language.FederatedType(MODEL_TYPE, federated_language.SERVER)\n",
- "CLIENT_DATA_TYPE = federated_language.FederatedType(LOCAL_DATA_TYPE, federated_language.CLIENTS)"
+ "SERVER_MODEL_TYPE = federated_language.FederatedType(\n",
+ " MODEL_TYPE, federated_language.SERVER\n",
+ ")\n",
+ "CLIENT_DATA_TYPE = federated_language.FederatedType(\n",
+ " LOCAL_DATA_TYPE, federated_language.CLIENTS\n",
+ ")"
]
},
{
@@ -1066,10 +1091,13 @@
},
"outputs": [],
"source": [
- "@tff.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)\n",
+ "@federated_language.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)\n",
"def federated_eval(model, data):\n",
" return federated_language.federated_mean(\n",
- " federated_language.federated_map(local_eval, [federated_language.federated_broadcast(model), data]))"
+ " federated_language.federated_map(\n",
+ " local_eval, [federated_language.federated_broadcast(model), data]\n",
+ " )\n",
+ " )"
]
},
{
@@ -1147,10 +1175,13 @@
}
],
"source": [
- "print('initial_model loss =', federated_eval(initial_model,\n",
- " federated_train_data))\n",
- "print('locally_trained_model loss =',\n",
- " federated_eval(locally_trained_model, federated_train_data))"
+ "print(\n",
+ " 'initial_model loss =', federated_eval(initial_model, federated_train_data)\n",
+ ")\n",
+ "print(\n",
+ " 'locally_trained_model loss =',\n",
+ " federated_eval(locally_trained_model, federated_train_data),\n",
+ ")"
]
},
{
@@ -1184,17 +1215,25 @@
},
"outputs": [],
"source": [
- "SERVER_FLOAT_TYPE = federated_language.FederatedType(np.float32, federated_language.SERVER)\n",
+ "SERVER_FLOAT_TYPE = federated_language.FederatedType(\n",
+ " np.float32, federated_language.SERVER\n",
+ ")\n",
"\n",
"\n",
- "@tff.federated_computation(SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE,\n",
- " CLIENT_DATA_TYPE)\n",
+ "@federated_language.federated_computation(\n",
+ " SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE, CLIENT_DATA_TYPE\n",
+ ")\n",
"def federated_train(model, learning_rate, data):\n",
" return federated_language.federated_mean(\n",
- " federated_language.federated_map(local_train, [\n",
- " federated_language.federated_broadcast(model),\n",
- " federated_language.federated_broadcast(learning_rate), data\n",
- " ]))"
+ " federated_language.federated_map(\n",
+ " local_train,\n",
+ " [\n",
+ " federated_language.federated_broadcast(model),\n",
+ " federated_language.federated_broadcast(learning_rate),\n",
+ " data,\n",
+ " ],\n",
+ " )\n",
+ " )"
]
},
{
@@ -1270,8 +1309,10 @@
}
],
"source": [
- "print('initial_model test loss =',\n",
- " federated_eval(initial_model, federated_test_data))\n",
+ "print(\n",
+ " 'initial_model test loss =',\n",
+ " federated_eval(initial_model, federated_test_data),\n",
+ ")\n",
"print('trained_model test loss =', federated_eval(model, federated_test_data))"
]
},
diff --git a/docs/tutorials/federated_learning_for_image_classification.ipynb b/docs/tutorials/federated_learning_for_image_classification.ipynb
index 6c5561225..bf2cba700 100644
--- a/docs/tutorials/federated_learning_for_image_classification.ipynb
+++ b/docs/tutorials/federated_learning_for_image_classification.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -109,7 +109,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -160,7 +160,7 @@
"\n",
"np.random.seed(0)\n",
"\n",
- "tff.federated_computation(lambda: 'Hello, World!')()"
+ "tff.tensorflow_computation(lambda: 'Hello, World!')()"
]
},
{
@@ -274,7 +274,8 @@
],
"source": [
"example_dataset = emnist_train.create_tf_dataset_for_client(\n",
- " emnist_train.client_ids[0])\n",
+ " emnist_train.client_ids[0]\n",
+ ")\n",
"\n",
"example_element = next(iter(example_dataset))\n",
"\n",
@@ -290,6 +291,7 @@
"outputs": [],
"source": [
"from matplotlib import pyplot as plt\n",
+ "\n",
"plt.imshow(example_element['pixels'].numpy(), cmap='gray', aspect='equal')\n",
"plt.grid(False)\n",
"_ = plt.show()"
@@ -328,7 +330,7 @@
"j = 0\n",
"\n",
"for example in example_dataset.take(40):\n",
- " plt.subplot(4, 10, j+1)\n",
+ " plt.subplot(4, 10, j + 1)\n",
" plt.imshow(example['pixels'].numpy(), cmap='gray', aspect='equal')\n",
" plt.axis('off')\n",
" j += 1"
@@ -356,20 +358,20 @@
"f.suptitle('Label Counts for a Sample of Clients')\n",
"for i in range(6):\n",
" client_dataset = emnist_train.create_tf_dataset_for_client(\n",
- " emnist_train.client_ids[i])\n",
+ " emnist_train.client_ids[i]\n",
+ " )\n",
" plot_data = collections.defaultdict(list)\n",
" for example in client_dataset:\n",
" # Append counts individually per label to make plots\n",
" # more colorful instead of one color per plot.\n",
" label = example['label'].numpy()\n",
" plot_data[label].append(label)\n",
- " plt.subplot(2, 3, i+1)\n",
+ " plt.subplot(2, 3, i + 1)\n",
" plt.title('Client {}'.format(i))\n",
" for j in range(10):\n",
" plt.hist(\n",
- " plot_data[j],\n",
- " density=False,\n",
- " bins=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])"
+ " plot_data[j], density=False, bins=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]\n",
+ " )"
]
},
{
@@ -394,7 +396,8 @@
"\n",
"for i in range(5):\n",
" client_dataset = emnist_train.create_tf_dataset_for_client(\n",
- " emnist_train.client_ids[i])\n",
+ " emnist_train.client_ids[i]\n",
+ " )\n",
" plot_data = collections.defaultdict(list)\n",
" for example in client_dataset:\n",
" plot_data[example['label'].numpy()].append(example['pixels'].numpy())\n",
@@ -402,7 +405,7 @@
" f.suptitle(\"Client #{}'s Mean Image Per Label\".format(i))\n",
" for j in range(10):\n",
" mean_img = np.mean(plot_data[j], 0)\n",
- " plt.subplot(2, 5, j+1)\n",
+ " plt.subplot(2, 5, j + 1)\n",
" plt.imshow(mean_img.reshape((28, 28)))\n",
" plt.axis('off')"
]
@@ -451,16 +454,23 @@
"SHUFFLE_BUFFER = 100\n",
"PREFETCH_BUFFER = 10\n",
"\n",
+ "\n",
"def preprocess(dataset):\n",
"\n",
" def batch_format_fn(element):\n",
" \"\"\"Flatten a batch `pixels` and return the features as an `OrderedDict`.\"\"\"\n",
" return collections.OrderedDict(\n",
" x=tf.reshape(element['pixels'], [-1, 784]),\n",
- " y=tf.reshape(element['label'], [-1, 1]))\n",
+ " y=tf.reshape(element['label'], [-1, 1]),\n",
+ " )\n",
"\n",
- " return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER, seed=1).batch(\n",
- " BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)"
+ " return (\n",
+ " dataset.repeat(NUM_EPOCHS)\n",
+ " .shuffle(SHUFFLE_BUFFER, seed=1)\n",
+ " .batch(BATCH_SIZE)\n",
+ " .map(batch_format_fn)\n",
+ " .prefetch(PREFETCH_BUFFER)\n",
+ " )"
]
},
{
@@ -518,8 +528,9 @@
"source": [
"preprocessed_example_dataset = preprocess(example_dataset)\n",
"\n",
- "sample_batch = tf.nest.map_structure(lambda x: x.numpy(),\n",
- " next(iter(preprocessed_example_dataset)))\n",
+ "sample_batch = tf.nest.map_structure(\n",
+ " lambda x: x.numpy(), next(iter(preprocessed_example_dataset))\n",
+ ")\n",
"\n",
"sample_batch"
]
@@ -676,7 +687,8 @@
" keras_model,\n",
" input_spec=preprocessed_example_dataset.element_spec,\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
- " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
+ " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],\n",
+ " )"
]
},
{
@@ -719,7 +731,8 @@
"training_process = tff.learning.algorithms.build_weighted_fed_avg(\n",
" model_fn,\n",
" client_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=0.02),\n",
- " server_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=1.0))"
+ " server_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=1.0),\n",
+ ")"
]
},
{
@@ -939,12 +952,12 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"logdir = \"/tmp/logs/scalars/training/\"\n",
"try:\n",
" tf.io.gfile.rmtree(logdir) # delete any previous results\n",
"except tf.errors.NotFoundError as e:\n",
- " pass # Ignore if the directory didn't previously exist.\n",
+ " pass # Ignore if the directory didn't previously exist.\n",
"summary_writer = tf.summary.create_file_writer(logdir)\n",
"train_state = training_process.initialize()"
]
@@ -966,7 +979,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"with summary_writer.as_default():\n",
" for round_num in range(1, NUM_ROUNDS):\n",
" result = training_process.next(train_state, federated_train_data)\n",
@@ -993,7 +1006,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!ls {logdir}\n",
"%tensorboard --logdir {logdir} --port=0"
]
@@ -1006,9 +1019,9 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"# Uncomment and run this cell to clean your directory of old output for\n",
- "# future graphs from this directory. We don't run it by default so that if \n",
+ "# future graphs from this directory. We don't run it by default so that if\n",
"# you do a \"Runtime \u003e Run all\" you don't lose your results.\n",
"\n",
"# !rm -R /tmp/logs/scalars/*"
@@ -1057,7 +1070,8 @@
"outputs": [],
"source": [
"MnistVariables = collections.namedtuple(\n",
- " 'MnistVariables', 'weights bias num_examples loss_sum accuracy_sum')"
+ " 'MnistVariables', 'weights bias num_examples loss_sum accuracy_sum'\n",
+ ")"
]
},
{
@@ -1086,14 +1100,17 @@
" weights=tf.Variable(\n",
" lambda: tf.zeros(dtype=tf.float32, shape=(784, 10)),\n",
" name='weights',\n",
- " trainable=True),\n",
+ " trainable=True,\n",
+ " ),\n",
" bias=tf.Variable(\n",
" lambda: tf.zeros(dtype=tf.float32, shape=(10)),\n",
" name='bias',\n",
- " trainable=True),\n",
+ " trainable=True,\n",
+ " ),\n",
" num_examples=tf.Variable(0.0, name='num_examples', trainable=False),\n",
" loss_sum=tf.Variable(0.0, name='loss_sum', trainable=False),\n",
- " accuracy_sum=tf.Variable(0.0, name='accuracy_sum', trainable=False))"
+ " accuracy_sum=tf.Variable(0.0, name='accuracy_sum', trainable=False),\n",
+ " )"
]
},
{
@@ -1119,15 +1136,18 @@
"def predict_on_batch(variables, x):\n",
" return tf.nn.softmax(tf.matmul(x, variables.weights) + variables.bias)\n",
"\n",
+ "\n",
"def mnist_forward_pass(variables, batch):\n",
" y = predict_on_batch(variables, batch['x'])\n",
" predictions = tf.cast(tf.argmax(y, 1), tf.int32)\n",
"\n",
" flat_labels = tf.reshape(batch['y'], [-1])\n",
" loss = -tf.reduce_mean(\n",
- " tf.reduce_sum(tf.one_hot(flat_labels, 10) * tf.math.log(y), axis=[1]))\n",
+ " tf.reduce_sum(tf.one_hot(flat_labels, 10) * tf.math.log(y), axis=[1])\n",
+ " )\n",
" accuracy = tf.reduce_mean(\n",
- " tf.cast(tf.equal(predictions, flat_labels), tf.float32))\n",
+ " tf.cast(tf.equal(predictions, flat_labels), tf.float32)\n",
+ " )\n",
"\n",
" num_examples = tf.cast(tf.size(batch['y']), tf.float32)\n",
"\n",
@@ -1161,7 +1181,8 @@
" return collections.OrderedDict(\n",
" num_examples=[variables.num_examples],\n",
" loss=[variables.loss_sum, variables.num_examples],\n",
- " accuracy=[variables.accuracy_sum, variables.num_examples])"
+ " accuracy=[variables.accuracy_sum, variables.num_examples],\n",
+ " )"
]
},
{
@@ -1185,7 +1206,8 @@
" return collections.OrderedDict(\n",
" num_examples=tf.function(func=lambda x: x[0]),\n",
" loss=tf.function(func=lambda x: x[0] / x[1]),\n",
- " accuracy=tf.function(func=lambda x: x[0] / x[1]))"
+ " accuracy=tf.function(func=lambda x: x[0] / x[1]),\n",
+ " )"
]
},
{
@@ -1221,6 +1243,7 @@
"import collections\n",
"from collections.abc import Callable\n",
"\n",
+ "\n",
"class MnistModel(tff.learning.models.VariableModel):\n",
"\n",
" def __init__(self):\n",
@@ -1237,37 +1260,42 @@
" @property\n",
" def local_variables(self):\n",
" return [\n",
- " self._variables.num_examples, self._variables.loss_sum,\n",
- " self._variables.accuracy_sum\n",
+ " self._variables.num_examples,\n",
+ " self._variables.loss_sum,\n",
+ " self._variables.accuracy_sum,\n",
" ]\n",
"\n",
" @property\n",
" def input_spec(self):\n",
" return collections.OrderedDict(\n",
" x=tf.TensorSpec([None, 784], tf.float32),\n",
- " y=tf.TensorSpec([None, 1], tf.int32))\n",
+ " y=tf.TensorSpec([None, 1], tf.int32),\n",
+ " )\n",
"\n",
" @tf.function\n",
" def predict_on_batch(self, x, training=True):\n",
" del training\n",
" return predict_on_batch(self._variables, x)\n",
- " \n",
+ "\n",
" @tf.function\n",
" def forward_pass(self, batch, training=True):\n",
" del training\n",
" loss, predictions = mnist_forward_pass(self._variables, batch)\n",
" num_exmaples = tf.shape(batch['x'])[0]\n",
" return tff.learning.models.BatchOutput(\n",
- " loss=loss, predictions=predictions, num_examples=num_exmaples)\n",
+ " loss=loss, predictions=predictions, num_examples=num_exmaples\n",
+ " )\n",
"\n",
" @tf.function\n",
" def report_local_unfinalized_metrics(\n",
- " self) -\u003e collections.OrderedDict[str, list[tf.Tensor]]:\n",
+ " self,\n",
+ " ) -\u003e collections.OrderedDict[str, list[tf.Tensor]]:\n",
" \"\"\"Creates an `OrderedDict` of metric names to unfinalized values.\"\"\"\n",
" return get_local_unfinalized_metrics(self._variables)\n",
"\n",
" def metric_finalizers(\n",
- " self) -\u003e collections.OrderedDict[str, Callable[[list[tf.Tensor]], tf.Tensor]]:\n",
+ " self,\n",
+ " ) -\u003e collections.OrderedDict[str, Callable[[list[tf.Tensor]], tf.Tensor]]:\n",
" \"\"\"Creates an `OrderedDict` of metric names to finalizers.\"\"\"\n",
" return get_metric_finalizers()\n",
"\n",
@@ -1340,7 +1368,8 @@
"source": [
"training_process = tff.learning.algorithms.build_weighted_fed_avg(\n",
" MnistModel,\n",
- " client_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=0.02))"
+ " client_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=0.02),\n",
+ ")"
]
},
{
@@ -1598,7 +1627,9 @@
"source": [
"evaluation_state = evaluation_process.initialize()\n",
"model_weights = training_process.get_model_weights(train_state)\n",
- "evaluation_state = evaluation_process.set_model_weights(evaluation_state, model_weights)"
+ "evaluation_state = evaluation_process.set_model_weights(\n",
+ " evaluation_state, model_weights\n",
+ ")"
]
},
{
@@ -1620,7 +1651,9 @@
},
"outputs": [],
"source": [
- "evaluation_output = evaluation_process.next(evaluation_state, federated_train_data)"
+ "evaluation_output = evaluation_process.next(\n",
+ " evaluation_state, federated_train_data\n",
+ ")"
]
},
{
@@ -1705,7 +1738,9 @@
},
"outputs": [],
"source": [
- "evaluation_output = evaluation_process.next(evaluation_state, federated_test_data)"
+ "evaluation_output = evaluation_process.next(\n",
+ " evaluation_state, federated_test_data\n",
+ ")"
]
},
{
diff --git a/docs/tutorials/federated_learning_for_text_generation.ipynb b/docs/tutorials/federated_learning_for_text_generation.ipynb
index f98e22ad7..1567b6501 100644
--- a/docs/tutorials/federated_learning_for_text_generation.ipynb
+++ b/docs/tutorials/federated_learning_for_text_generation.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -85,7 +85,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
},
@@ -120,7 +120,7 @@
"np.random.seed(0)\n",
"\n",
"# Test that TFF is working:\n",
- "tff.federated_computation(lambda: 'Hello, World!')()"
+ "tff.tensorflow_computation(lambda: 'Hello, World!')()"
]
},
{
@@ -161,10 +161,13 @@
"outputs": [],
"source": [
"# A fixed vocabularly of ASCII chars that occur in the works of Shakespeare and Dickens:\n",
- "vocab = list('dhlptx@DHLPTX $(,048cgkoswCGKOSW[_#\\'/37;?bfjnrvzBFJNRVZ\"\u0026*.26:\\naeimquyAEIMQUY]!%)-159\\r')\n",
+ "vocab = list(\n",
+ " 'dhlptx@DHLPTX'\n",
+ " ' $(,048cgkoswCGKOSW[_#\\'/37;?bfjnrvzBFJNRVZ\"\u0026*.26:\\naeimquyAEIMQUY]!%)-159\\r'\n",
+ ")\n",
"\n",
"# Creating a mapping from unique characters to indices\n",
- "char2idx = {u:i for i, u in enumerate(vocab)}\n",
+ "char2idx = {u: i for i, u in enumerate(vocab)}\n",
"idx2char = np.array(vocab)"
]
},
@@ -188,10 +191,11 @@
"def load_model(batch_size):\n",
" urls = {\n",
" 1: 'https://storage.googleapis.com/tff-models-public/dickens_rnn.batch1.kerasmodel',\n",
- " 8: 'https://storage.googleapis.com/tff-models-public/dickens_rnn.batch8.kerasmodel'}\n",
+ " 8: 'https://storage.googleapis.com/tff-models-public/dickens_rnn.batch8.kerasmodel',\n",
+ " }\n",
" assert batch_size in urls, 'batch_size must be in ' + str(urls.keys())\n",
" url = urls[batch_size]\n",
- " local_file = tf.keras.utils.get_file(os.path.basename(url), origin=url) \n",
+ " local_file = tf.keras.utils.get_file(os.path.basename(url), origin=url)\n",
" return tf.keras.models.load_model(local_file, compile=False)"
]
},
@@ -216,12 +220,13 @@
" predictions = model(input_eval)\n",
" predictions = tf.squeeze(predictions, 0)\n",
" predictions = predictions / temperature\n",
- " predicted_id = tf.random.categorical(\n",
- " predictions, num_samples=1)[-1, 0].numpy()\n",
+ " predicted_id = tf.random.categorical(predictions, num_samples=1)[\n",
+ " -1, 0\n",
+ " ].numpy()\n",
" input_eval = tf.expand_dims([predicted_id], 0)\n",
" text_generated.append(idx2char[predicted_id])\n",
"\n",
- " return (start_string + ''.join(text_generated))"
+ " return start_string + ''.join(text_generated)"
]
},
{
@@ -247,7 +252,9 @@
"source": [
"# Text generation requires a batch_size=1 model.\n",
"keras_model_batch1 = load_model(batch_size=1)\n",
- "print(generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? '))"
+ "print(\n",
+ " generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? ')\n",
+ ")"
]
},
{
@@ -309,7 +316,8 @@
"source": [
"# Here the play is \"The Tragedy of King Lear\" and the character is \"King\".\n",
"raw_example_dataset = train_data.create_tf_dataset_for_client(\n",
- " 'THE_TRAGEDY_OF_KING_LEAR_KING')\n",
+ " 'THE_TRAGEDY_OF_KING_LEAR_KING'\n",
+ ")\n",
"# To allow for future extensions, each entry x\n",
"# is an OrderedDict with a single key 'snippets' which contains the text.\n",
"for x in raw_example_dataset.take(2):\n",
@@ -351,9 +359,10 @@
"# using the vocab loaded above:\n",
"table = tf.lookup.StaticHashTable(\n",
" tf.lookup.KeyValueTensorInitializer(\n",
- " keys=vocab, values=tf.constant(list(range(len(vocab))),\n",
- " dtype=tf.int64)),\n",
- " default_value=0)\n",
+ " keys=vocab, values=tf.constant(list(range(len(vocab))), dtype=tf.int64)\n",
+ " ),\n",
+ " default_value=0,\n",
+ ")\n",
"\n",
"\n",
"def to_ids(x):\n",
@@ -381,7 +390,8 @@
" .shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)\n",
" # And finally split into (input, target) tuples,\n",
" # each of length SEQ_LENGTH.\n",
- " .map(split_input_target))"
+ " .map(split_input_target)\n",
+ " )"
]
},
{
@@ -506,27 +516,35 @@
}
],
"source": [
- "BATCH_SIZE = 8 # The training and eval batch size for the rest of this tutorial.\n",
+ "BATCH_SIZE = (\n",
+ " 8 # The training and eval batch size for the rest of this tutorial.\n",
+ ")\n",
"keras_model = load_model(batch_size=BATCH_SIZE)\n",
"keras_model.compile(\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
- " metrics=[FlattenedCategoricalAccuracy()])\n",
+ " metrics=[FlattenedCategoricalAccuracy()],\n",
+ ")\n",
"\n",
"# Confirm that loss is much lower on Shakespeare than on random data\n",
"loss, accuracy = keras_model.evaluate(example_dataset.take(5), verbose=0)\n",
"print(\n",
- " 'Evaluating on an example Shakespeare character: {a:3f}'.format(a=accuracy))\n",
+ " 'Evaluating on an example Shakespeare character: {a:3f}'.format(a=accuracy)\n",
+ ")\n",
"\n",
"# As a sanity check, we can construct some completely random data, where we expect\n",
"# the accuracy to be essentially random:\n",
"random_guessed_accuracy = 1.0 / len(vocab)\n",
- "print('Expected accuracy for random guessing: {a:.3f}'.format(\n",
- " a=random_guessed_accuracy))\n",
+ "print(\n",
+ " 'Expected accuracy for random guessing: {a:.3f}'.format(\n",
+ " a=random_guessed_accuracy\n",
+ " )\n",
+ ")\n",
"random_indexes = np.random.randint(\n",
- " low=0, high=len(vocab), size=1 * BATCH_SIZE * (SEQ_LENGTH + 1))\n",
+ " low=0, high=len(vocab), size=1 * BATCH_SIZE * (SEQ_LENGTH + 1)\n",
+ ")\n",
"data = collections.OrderedDict(\n",
- " snippets=tf.constant(\n",
- " ''.join(np.array(vocab)[random_indexes]), shape=[1, 1]))\n",
+ " snippets=tf.constant(''.join(np.array(vocab)[random_indexes]), shape=[1, 1])\n",
+ ")\n",
"random_dataset = preprocess(tf.data.Dataset.from_tensor_slices(data))\n",
"loss, accuracy = keras_model.evaluate(random_dataset, steps=10, verbose=0)\n",
"print('Evaluating on completely random data: {a:.3f}'.format(a=accuracy))"
@@ -563,8 +581,8 @@
"outputs": [],
"source": [
"# Clone the keras_model inside `create_tff_model()`, which TFF will\n",
- "# call to produce a new copy of the model inside the graph that it will \n",
- "# serialize. Note: we want to construct all the necessary objects we'll need \n",
+ "# call to produce a new copy of the model inside the graph that it will\n",
+ "# serialize. Note: we want to construct all the necessary objects we'll need\n",
"# _inside_ this method.\n",
"def create_tff_model():\n",
" # TFF uses an `input_spec` so it knows the types and shapes\n",
@@ -575,7 +593,8 @@
" keras_model_clone,\n",
" input_spec=input_spec,\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
- " metrics=[FlattenedCategoricalAccuracy()])"
+ " metrics=[FlattenedCategoricalAccuracy()],\n",
+ " )"
]
},
{
@@ -599,10 +618,11 @@
},
"outputs": [],
"source": [
- "# This command builds all the TensorFlow graphs and serializes them: \n",
+ "# This command builds all the TensorFlow graphs and serializes them:\n",
"fed_avg = tff.learning.algorithms.build_weighted_fed_avg(\n",
" model_fn=create_tff_model,\n",
- " client_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=0.5))"
+ " client_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=0.5),\n",
+ ")"
]
},
{
@@ -634,8 +654,11 @@
"result = fed_avg.next(state, [example_dataset.take(5)])\n",
"state = result.state\n",
"train_metrics = result.metrics['client_work']['train']\n",
- "print('loss={l:.3f}, accuracy={a:.3f}'.format(\n",
- " l=train_metrics['loss'], a=train_metrics['accuracy']))"
+ "print(\n",
+ " 'loss={l:.3f}, accuracy={a:.3f}'.format(\n",
+ " l=train_metrics['loss'], a=train_metrics['accuracy']\n",
+ " )\n",
+ ")"
]
},
{
@@ -662,15 +685,17 @@
"\n",
"\n",
"clients = [\n",
- " 'ALL_S_WELL_THAT_ENDS_WELL_CELIA', 'MUCH_ADO_ABOUT_NOTHING_OTHELLO',\n",
+ " 'ALL_S_WELL_THAT_ENDS_WELL_CELIA',\n",
+ " 'MUCH_ADO_ABOUT_NOTHING_OTHELLO',\n",
"]\n",
"\n",
"train_datasets = [data(client) for client in clients]\n",
"\n",
- "# We concatenate the test datasets for evaluation with Keras by creating a \n",
+ "# We concatenate the test datasets for evaluation with Keras by creating a\n",
"# Dataset of Datasets, and then identity flat mapping across all the examples.\n",
"test_dataset = tf.data.Dataset.from_tensor_slices(\n",
- " [data(client, test_data) for client in clients]).flat_map(lambda x: x)"
+ " [data(client, test_data) for client in clients]\n",
+ ").flat_map(lambda x: x)"
]
},
{
@@ -726,7 +751,7 @@
"# Load our pre-trained Keras model weights into the global model state.\n",
"pre_trained_weights = tff.learning.models.ModelWeights(\n",
" trainable=[v.numpy() for v in keras_model.trainable_weights],\n",
- " non_trainable=[v.numpy() for v in keras_model.non_trainable_weights]\n",
+ " non_trainable=[v.numpy() for v in keras_model.non_trainable_weights],\n",
")\n",
"state = fed_avg.set_model_weights(state, pre_trained_weights)\n",
"\n",
@@ -737,7 +762,8 @@
" keras_model = load_model(batch_size=BATCH_SIZE)\n",
" keras_model.compile(\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
- " metrics=[FlattenedCategoricalAccuracy()])\n",
+ " metrics=[FlattenedCategoricalAccuracy()],\n",
+ " )\n",
" model_weights = fed_avg.get_model_weights(state)\n",
" model_weights.assign_weights_to(keras_model)\n",
" loss, accuracy = keras_model.evaluate(example_dataset, steps=2, verbose=0)\n",
@@ -750,8 +776,11 @@
" result = fed_avg.next(state, train_datasets)\n",
" state = result.state\n",
" train_metrics = result.metrics['client_work']['train']\n",
- " print('\\tTrain: loss={l:.3f}, accuracy={a:.3f}'.format(\n",
- " l=train_metrics['loss'], a=train_metrics['accuracy']))\n",
+ " print(\n",
+ " '\\tTrain: loss={l:.3f}, accuracy={a:.3f}'.format(\n",
+ " l=train_metrics['loss'], a=train_metrics['accuracy']\n",
+ " )\n",
+ " )\n",
"\n",
"print('Final evaluation')\n",
"keras_evaluate(state, NUM_ROUNDS + 1)"
@@ -788,7 +817,9 @@
"# Set our newly trained weights back in the originally created model.\n",
"keras_model_batch1.set_weights([v.numpy() for v in keras_model.weights])\n",
"# Text generation requires batch_size=1\n",
- "print(generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? '))"
+ "print(\n",
+ " generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? ')\n",
+ ")"
]
},
{
diff --git a/docs/tutorials/federated_learning_with_differential_privacy.ipynb b/docs/tutorials/federated_learning_with_differential_privacy.ipynb
index 14cf821af..39f822383 100644
--- a/docs/tutorials/federated_learning_with_differential_privacy.ipynb
+++ b/docs/tutorials/federated_learning_with_differential_privacy.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -87,7 +87,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quiet --upgrade dp-accounting\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -150,10 +150,11 @@
}
],
"source": [
- "@tff.federated_computation\n",
+ "@tff.tensorflow_computation\n",
"def hello_world():\n",
" return 'Hello, World!'\n",
"\n",
+ "\n",
"hello_world()"
]
},
@@ -176,28 +177,34 @@
"source": [
"def get_emnist_dataset():\n",
" emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(\n",
- " only_digits=True)\n",
+ " only_digits=True\n",
+ " )\n",
"\n",
" def element_fn(element):\n",
" return collections.OrderedDict(\n",
- " x=tf.expand_dims(element['pixels'], -1), y=element['label'])\n",
+ " x=tf.expand_dims(element['pixels'], -1), y=element['label']\n",
+ " )\n",
"\n",
" def preprocess_train_dataset(dataset):\n",
" # Use buffer_size same as the maximum client dataset size,\n",
" # 418 for Federated EMNIST\n",
- " return (dataset.map(element_fn)\n",
- " .shuffle(buffer_size=418)\n",
- " .repeat(1)\n",
- " .batch(32, drop_remainder=False))\n",
+ " return (\n",
+ " dataset.map(element_fn)\n",
+ " .shuffle(buffer_size=418)\n",
+ " .repeat(1)\n",
+ " .batch(32, drop_remainder=False)\n",
+ " )\n",
"\n",
" def preprocess_test_dataset(dataset):\n",
" return dataset.map(element_fn).batch(128, drop_remainder=False)\n",
"\n",
" emnist_train = emnist_train.preprocess(preprocess_train_dataset)\n",
" emnist_test = preprocess_test_dataset(\n",
- " emnist_test.create_tf_dataset_from_all_clients())\n",
+ " emnist_test.create_tf_dataset_from_all_clients()\n",
+ " )\n",
" return emnist_train, emnist_test\n",
"\n",
+ "\n",
"train_data, test_data = get_emnist_dataset()"
]
},
@@ -223,12 +230,14 @@
" tf.keras.layers.Reshape(input_shape=(28, 28, 1), target_shape=(28 * 28,)),\n",
" tf.keras.layers.Dense(200, activation=tf.nn.relu),\n",
" tf.keras.layers.Dense(200, activation=tf.nn.relu),\n",
- " tf.keras.layers.Dense(10)])\n",
+ " tf.keras.layers.Dense(10),\n",
+ " ])\n",
" return tff.learning.models.from_keras_model(\n",
" keras_model=model,\n",
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
" input_spec=test_data.element_spec,\n",
- " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
+ " metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],\n",
+ " )"
]
},
{
@@ -258,11 +267,13 @@
"source": [
"total_clients = len(train_data.client_ids)\n",
"\n",
+ "\n",
"def train(rounds, noise_multiplier, clients_per_round, data_frame):\n",
" # Using the `dp_aggregator` here turns on differential privacy with adaptive\n",
" # clipping.\n",
" aggregation_factory = tff.learning.model_update_aggregator.dp_aggregator(\n",
- " noise_multiplier, clients_per_round)\n",
+ " noise_multiplier, clients_per_round\n",
+ " )\n",
"\n",
" # We use Poisson subsampling which gives slightly tighter privacy guarantees\n",
" # compared to having a fixed number of clients per round. The actual number of\n",
@@ -274,14 +285,14 @@
" # updates can cause the second moment accumulators to become very large\n",
" # prematurely.\n",
" learning_process = tff.learning.algorithms.build_unweighted_fed_avg(\n",
- " my_model_fn,\n",
- " client_optimizer_fn=tff.learning.optimizers.build_sgdm(0.01),\n",
- " server_optimizer_fn=tff.learning.optimizers.build_sgdm(\n",
- " 1.0, momentum=0.9),\n",
- " model_aggregator=aggregation_factory)\n",
+ " my_model_fn,\n",
+ " client_optimizer_fn=tff.learning.optimizers.build_sgdm(0.01),\n",
+ " server_optimizer_fn=tff.learning.optimizers.build_sgdm(1.0, momentum=0.9),\n",
+ " model_aggregator=aggregation_factory,\n",
+ " )\n",
"\n",
" eval_process = tff.learning.build_federated_evaluation(my_model_fn)\n",
- " \n",
+ "\n",
" # Training loop.\n",
" state = learning_process.initialize()\n",
" for round in range(rounds):\n",
@@ -290,19 +301,23 @@
" metrics = eval_process(model_weights, [test_data])['eval']\n",
" if round \u003c 25 or round % 25 == 0:\n",
" print(f'Round {round:3d}: {metrics}')\n",
- " data_frame = data_frame.append({'Round': round,\n",
- " 'NoiseMultiplier': noise_multiplier,\n",
- " **metrics}, ignore_index=True)\n",
+ " data_frame = data_frame.append(\n",
+ " {'Round': round, 'NoiseMultiplier': noise_multiplier, **metrics},\n",
+ " ignore_index=True,\n",
+ " )\n",
"\n",
" # Sample clients for a round. Note that if your dataset is large and\n",
" # sampling_prob is small, it would be faster to use gap sampling.\n",
" x = np.random.uniform(size=total_clients)\n",
" sampled_clients = [\n",
- " train_data.client_ids[i] for i in range(total_clients)\n",
- " if x[i] \u003c sampling_prob]\n",
+ " train_data.client_ids[i]\n",
+ " for i in range(total_clients)\n",
+ " if x[i] \u003c sampling_prob\n",
+ " ]\n",
" sampled_train_data = [\n",
" train_data.create_tf_dataset_for_client(client)\n",
- " for client in sampled_clients]\n",
+ " for client in sampled_clients\n",
+ " ]\n",
"\n",
" # Use selected clients for update.\n",
" result = learning_process.next(state, sampled_train_data)\n",
@@ -312,11 +327,12 @@
" model_weights = learning_process.get_model_weights(state)\n",
" metrics = eval_process(model_weights, [test_data])['eval']\n",
" print(f'Round {rounds:3d}: {metrics}')\n",
- " data_frame = data_frame.append({'Round': rounds,\n",
- " 'NoiseMultiplier': noise_multiplier,\n",
- " **metrics}, ignore_index=True)\n",
+ " data_frame = data_frame.append(\n",
+ " {'Round': rounds, 'NoiseMultiplier': noise_multiplier, **metrics},\n",
+ " ignore_index=True,\n",
+ " )\n",
"\n",
- " return data_frame\n"
+ " return data_frame"
]
},
{
@@ -408,16 +424,22 @@
"import matplotlib.pyplot as plt\n",
"import seaborn as sns\n",
"\n",
+ "\n",
"def make_plot(data_frame):\n",
" plt.figure(figsize=(15, 5))\n",
"\n",
" dff = data_frame.rename(\n",
- " columns={'sparse_categorical_accuracy': 'Accuracy', 'loss': 'Loss'})\n",
+ " columns={'sparse_categorical_accuracy': 'Accuracy', 'loss': 'Loss'}\n",
+ " )\n",
"\n",
" plt.subplot(121)\n",
- " sns.lineplot(data=dff, x='Round', y='Accuracy', hue='NoiseMultiplier', palette='dark')\n",
+ " sns.lineplot(\n",
+ " data=dff, x='Round', y='Accuracy', hue='NoiseMultiplier', palette='dark'\n",
+ " )\n",
" plt.subplot(122)\n",
- " sns.lineplot(data=dff, x='Round', y='Loss', hue='NoiseMultiplier', palette='dark')"
+ " sns.lineplot(\n",
+ " data=dff, x='Round', y='Loss', hue='NoiseMultiplier', palette='dark'\n",
+ " )"
]
},
{
@@ -487,7 +509,8 @@
"# No-arg callable that returns a fresh accountant.\n",
"make_fresh_accountant = dp_accounting.rdp.RdpAccountant\n",
"\n",
- "# Create function that takes expected clients per round and returns a \n",
+ "\n",
+ "# Create function that takes expected clients per round and returns a\n",
"# dp_accounting.DpEvent representing the full training process.\n",
"def make_event_from_param(clients_per_round):\n",
" q = clients_per_round / total_clients\n",
@@ -497,18 +520,25 @@
" composed_event = dp_accounting.SelfComposedDpEvent(sampled_event, rounds)\n",
" return composed_event\n",
"\n",
+ "\n",
"# Create object representing the search range [1, 3383].\n",
"bracket_interval = dp_accounting.ExplicitBracketInterval(1, total_clients)\n",
"\n",
"# Perform search for smallest clients_per_round achieving the target privacy.\n",
"clients_per_round = dp_accounting.calibrate_dp_mechanism(\n",
- " make_fresh_accountant, make_event_from_param, target_eps, target_delta,\n",
- " bracket_interval, discrete=True\n",
+ " make_fresh_accountant,\n",
+ " make_event_from_param,\n",
+ " target_eps,\n",
+ " target_delta,\n",
+ " bracket_interval,\n",
+ " discrete=True,\n",
")\n",
"\n",
"noise_multiplier = clients_per_round * noise_to_clients_ratio\n",
- "print(f'To get ({target_eps}, {target_delta})-DP, use {clients_per_round} '\n",
- " f'clients with noise multiplier {noise_multiplier}.')"
+ "print(\n",
+ " f'To get ({target_eps}, {target_delta})-DP, use {clients_per_round} '\n",
+ " f'clients with noise multiplier {noise_multiplier}.'\n",
+ ")"
]
},
{
diff --git a/docs/tutorials/federated_select.ipynb b/docs/tutorials/federated_select.ipynb
index 3435f6e1b..27c9eb253 100644
--- a/docs/tutorials/federated_select.ipynb
+++ b/docs/tutorials/federated_select.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -81,7 +81,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -135,6 +135,7 @@
"select_fn = tff.tensorflow.computation(lambda val, index: tf.gather(val, index))\n",
"client_data_type = np.str_\n",
"\n",
+ "\n",
"# A function from our client data to the indices of the values we'd like to\n",
"# select from the server.\n",
"@tff.tensorflow.computation(client_data_type)\n",
@@ -144,20 +145,40 @@
" split = tf.strings.split([client_string], sep=',')[0]\n",
" return tf.strings.to_number([split[0], split[1]], tf.int32)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation(federated_language.SequenceType(np.str_))\n",
"def concatenate(values):\n",
" def reduce_fn(acc, item):\n",
- " return tf.cond(tf.math.equal(acc, ''),\n",
- " lambda: item,\n",
- " lambda: tf.strings.join([acc, item], ','))\n",
+ " return tf.cond(\n",
+ " tf.math.equal(acc, ''),\n",
+ " lambda: item,\n",
+ " lambda: tf.strings.join([acc, item], ','),\n",
+ " )\n",
+ "\n",
" return values.reduce('', reduce_fn)\n",
"\n",
- "@tff.federated_computation(federated_language.FederatedType(list_of_strings_type, federated_language.SERVER), federated_language.FederatedType(client_data_type, federated_language.CLIENTS))\n",
+ "\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(\n",
+ " list_of_strings_type, federated_language.SERVER\n",
+ " ),\n",
+ " federated_language.FederatedType(\n",
+ " client_data_type, federated_language.CLIENTS\n",
+ " ),\n",
+ ")\n",
"def broadcast_based_on_client_data(list_of_strings_at_server, client_data):\n",
- " keys_at_clients = federated_language.federated_map(keys_for_client, client_data)\n",
- " max_key = federated_language.federated_map(get_size, list_of_strings_at_server)\n",
- " values_at_clients = federated_language.federated_select(keys_at_clients, max_key, list_of_strings_at_server, select_fn)\n",
- " value_at_clients = federated_language.federated_map(concatenate, values_at_clients)\n",
+ " keys_at_clients = federated_language.federated_map(\n",
+ " keys_for_client, client_data\n",
+ " )\n",
+ " max_key = federated_language.federated_map(\n",
+ " get_size, list_of_strings_at_server\n",
+ " )\n",
+ " values_at_clients = federated_language.federated_select(\n",
+ " keys_at_clients, max_key, list_of_strings_at_server, select_fn\n",
+ " )\n",
+ " value_at_clients = federated_language.federated_map(\n",
+ " concatenate, values_at_clients\n",
+ " )\n",
" return value_at_clients"
]
},
@@ -228,24 +249,38 @@
"def get_random_key(max_key):\n",
" return tf.random.uniform(shape=[1], minval=0, maxval=max_key, dtype=tf.int32)\n",
"\n",
+ "\n",
"list_of_strings_type = federated_language.TensorType(np.str_, [None])\n",
"get_size = tff.tensorflow.computation(lambda x: tf.size(x))\n",
"select_fn = tff.tensorflow.computation(lambda val, index: tf.gather(val, index))\n",
"\n",
+ "\n",
"@tff.tensorflow.computation(federated_language.SequenceType(np.str_))\n",
"def get_last_element(sequence):\n",
" return sequence.reduce('', lambda _initial_state, val: val)\n",
"\n",
- "@tff.federated_computation(federated_language.FederatedType(list_of_strings_type, federated_language.SERVER))\n",
+ "\n",
+ "@federated_language.federated_computation(\n",
+ " federated_language.FederatedType(\n",
+ " list_of_strings_type, federated_language.SERVER\n",
+ " )\n",
+ ")\n",
"def broadcast_random_element(list_of_strings_at_server):\n",
- " max_key_at_server = federated_language.federated_map(get_size, list_of_strings_at_server)\n",
+ " max_key_at_server = federated_language.federated_map(\n",
+ " get_size, list_of_strings_at_server\n",
+ " )\n",
" max_key_at_clients = federated_language.federated_broadcast(max_key_at_server)\n",
- " key_at_clients = federated_language.federated_map(get_random_key, max_key_at_clients)\n",
+ " key_at_clients = federated_language.federated_map(\n",
+ " get_random_key, max_key_at_clients\n",
+ " )\n",
" random_string_sequence_at_clients = federated_language.federated_select(\n",
- " key_at_clients, max_key_at_server, list_of_strings_at_server, select_fn)\n",
+ " key_at_clients, max_key_at_server, list_of_strings_at_server, select_fn\n",
+ " )\n",
" # Even though we only passed in a single key, `federated_select` returns a\n",
" # sequence for each client. We only care about the last (and only) element.\n",
- " random_string_at_clients = federated_language.federated_map(get_last_element, random_string_sequence_at_clients)\n",
+ " random_string_at_clients = federated_language.federated_map(\n",
+ " get_last_element, random_string_sequence_at_clients\n",
+ " )\n",
" return random_string_at_clients"
]
},
diff --git a/docs/tutorials/private_heavy_hitters.ipynb b/docs/tutorials/private_heavy_hitters.ipynb
index fbfb025d8..8abeff7b6 100644
--- a/docs/tutorials/private_heavy_hitters.ipynb
+++ b/docs/tutorials/private_heavy_hitters.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -94,7 +94,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"\n",
"# tensorflow_federated_nightly also bring in tf_nightly, which\n",
"# can causes a duplicate tensorboard install, leading to errors.\n",
@@ -124,6 +124,7 @@
"source": [
"import collections\n",
"\n",
+ "import federated_language\n",
"import numpy as np\n",
"import tensorflow as tf\n",
"import tensorflow_federated as tff\n",
@@ -132,7 +133,7 @@
"np.random.seed(0)\n",
"tff.backends.test.set_sync_test_cpp_execution_context()\n",
"\n",
- "tff.federated_computation(lambda: 'Hello, World!')()"
+ "federated_language.federated_computation(lambda: 'Hello, World!')()"
]
},
{
@@ -201,6 +202,7 @@
"# Preprocessing function to tokenize a line into words.\n",
"def tokenize(ds):\n",
" \"\"\"Tokenizes a line into words with alphanum characters.\"\"\"\n",
+ "\n",
" def extract_strings(example):\n",
" return tf.expand_dims(example['snippets'], 0)\n",
"\n",
@@ -209,7 +211,8 @@
"\n",
" def mask_all_symbolic_words(word):\n",
" return tf.math.logical_not(\n",
- " tf_text.wordshape(word, tf_text.WordShape.IS_PUNCT_OR_SYMBOL))\n",
+ " tf_text.wordshape(word, tf_text.WordShape.IS_PUNCT_OR_SYMBOL)\n",
+ " )\n",
"\n",
" tokenizer = tf_text.WhitespaceTokenizer()\n",
" ds = ds.map(extract_strings)\n",
@@ -218,11 +221,15 @@
" ds = ds.filter(mask_all_symbolic_words)\n",
" return ds\n",
"\n",
+ "\n",
"batch_size = 5\n",
"\n",
+ "\n",
"def client_data(n: int) -\u003e tf.data.Dataset:\n",
- " return tokenize(source.create_tf_dataset_for_client(\n",
- " source.client_ids[n])).batch(batch_size)\n",
+ " return tokenize(\n",
+ " source.create_tf_dataset_for_client(source.client_ids[n])\n",
+ " ).batch(batch_size)\n",
+ "\n",
"\n",
"# Pick a subset of client devices to participate in the computation.\n",
"dataset = [client_data(n) for n in range(10)]"
@@ -271,7 +278,8 @@
" max_heavy_hitters=10,\n",
" secure_sum_bitwidth=32,\n",
" multi_contribution=False,\n",
- " batch_size=batch_size)"
+ " batch_size=batch_size,\n",
+ ")"
]
},
{
@@ -296,7 +304,9 @@
},
"outputs": [],
"source": [
- "def run_simulation(one_round_computation: federated_language.Computation, dataset):\n",
+ "def run_simulation(\n",
+ " one_round_computation: federated_language.Computation, dataset\n",
+ "):\n",
" output = one_round_computation(dataset)\n",
" heavy_hitters = output.heavy_hitters\n",
" heavy_hitters_counts = output.heavy_hitters_counts\n",
@@ -372,7 +382,8 @@
" max_words_per_user=max_words_per_user,\n",
" secure_sum_bitwidth=32,\n",
" multi_contribution=False,\n",
- " batch_size=batch_size)\n",
+ " batch_size=batch_size,\n",
+ ")\n",
"\n",
"clients, result = run_simulation(iblt_computation, dataset)"
]
diff --git a/docs/tutorials/random_noise_generation.ipynb b/docs/tutorials/random_noise_generation.ipynb
index f42f851b8..34f54c520 100644
--- a/docs/tutorials/random_noise_generation.ipynb
+++ b/docs/tutorials/random_noise_generation.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -83,7 +83,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
@@ -133,10 +133,11 @@
}
],
"source": [
- "@tff.federated_computation\n",
+ "@federated_language.federated_computation\n",
"def hello_world():\n",
" return 'Hello, World!'\n",
"\n",
+ "\n",
"hello_world()"
]
},
@@ -187,25 +188,33 @@
"# Set to use 10 clients.\n",
"tff.backends.native.set_sync_local_cpp_execution_context(default_num_clients=10)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation\n",
"def noise_from_seed(seed):\n",
" return tf.random.stateless_normal((), seed=seed)\n",
"\n",
- "seed_type_at_server = federated_language.FederatedType(federated_language.to_type((np.int64, [2])), federated_language.SERVER)\n",
"\n",
- "@tff.federated_computation(seed_type_at_server)\n",
+ "seed_type_at_server = federated_language.FederatedType(\n",
+ " federated_language.to_type((np.int64, [2])), federated_language.SERVER\n",
+ ")\n",
+ "\n",
+ "\n",
+ "@federated_language.federated_computation(seed_type_at_server)\n",
"def get_random_min_and_max_deterministic(seed):\n",
" # Broadcast seed to all clients.\n",
" seed_on_clients = federated_language.federated_broadcast(seed)\n",
"\n",
" # Clients generate noise from seed deterministicly.\n",
- " noise_on_clients = federated_language.federated_map(noise_from_seed, seed_on_clients)\n",
+ " noise_on_clients = federated_language.federated_map(\n",
+ " noise_from_seed, seed_on_clients\n",
+ " )\n",
"\n",
" # Aggregate and return the min and max of the values generated on clients.\n",
" min = federated_language.federated_min(noise_on_clients)\n",
" max = federated_language.federated_max(noise_on_clients)\n",
" return min, max\n",
"\n",
+ "\n",
"seed = tf.constant([1, 1], dtype=tf.int64)\n",
"min, max = get_random_min_and_max_deterministic(seed)\n",
"assert min == max\n",
@@ -249,13 +258,17 @@
" gen = tf.random.Generator.from_non_deterministic_state()\n",
" return gen.normal(())\n",
"\n",
- "@tff.federated_computation\n",
+ "\n",
+ "@federated_language.federated_computation\n",
"def get_random_min_and_max_nondeterministic():\n",
- " noise_on_clients = federated_language.federated_eval(nondeterministic_noise, federated_language.CLIENTS)\n",
+ " noise_on_clients = federated_language.federated_eval(\n",
+ " nondeterministic_noise, federated_language.CLIENTS\n",
+ " )\n",
" min = federated_language.federated_min(noise_on_clients)\n",
" max = federated_language.federated_max(noise_on_clients)\n",
" return min, max\n",
"\n",
+ "\n",
"min, max = get_random_min_and_max_nondeterministic()\n",
"assert min != max\n",
"print(f'Values differ across clients. {min:8.3f},{max:8.3f}.')\n",
@@ -298,20 +311,26 @@
" outputs = tf.keras.layers.Dense(1)(inputs)\n",
" return tf.keras.Model(inputs=inputs, outputs=outputs)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation\n",
"def tff_return_model_init():\n",
" model = _keras_model()\n",
" # return the initialized single weight value of the dense layer\n",
" return tf.reshape(\n",
- " tff.learning.models.ModelWeights.from_model(model).trainable[0], [-1])[0]\n",
+ " tff.learning.models.ModelWeights.from_model(model).trainable[0], [-1]\n",
+ " )[0]\n",
+ "\n",
"\n",
- "@tff.federated_computation\n",
+ "@federated_language.federated_computation\n",
"def get_random_min_and_max_nondeterministic():\n",
- " noise_on_clients = federated_language.federated_eval(tff_return_model_init, federated_language.CLIENTS)\n",
+ " noise_on_clients = federated_language.federated_eval(\n",
+ " tff_return_model_init, federated_language.CLIENTS\n",
+ " )\n",
" min = federated_language.federated_min(noise_on_clients)\n",
" max = federated_language.federated_max(noise_on_clients)\n",
" return min, max\n",
"\n",
+ "\n",
"min, max = get_random_min_and_max_nondeterministic()\n",
"assert min != max\n",
"print(f'Values differ across clients. {min:8.3f},{max:8.3f}.')\n",
@@ -361,12 +380,14 @@
"source": [
"tf.random.set_seed(1)\n",
"\n",
+ "\n",
"@tf.function\n",
"def return_one_noise(_):\n",
" return tf.random.normal([])\n",
"\n",
- "n1=return_one_noise(1)\n",
- "n2=return_one_noise(2)\n",
+ "\n",
+ "n1 = return_one_noise(1)\n",
+ "n2 = return_one_noise(2)\n",
"assert n1 == n2\n",
"print(n1.numpy(), n2.numpy())"
]
@@ -399,12 +420,14 @@
"source": [
"tf.random.set_seed(1)\n",
"\n",
+ "\n",
"@tff.tensorflow.computation\n",
"def return_one_noise(_):\n",
" return tf.random.normal([])\n",
"\n",
- "n1=return_one_noise(1)\n",
- "n2=return_one_noise(2)\n",
+ "\n",
+ "n1 = return_one_noise(1)\n",
+ "n2 = return_one_noise(2)\n",
"assert n1 != n2\n",
"print(n1, n2)"
]
@@ -448,16 +471,20 @@
"source": [
"@tff.tensorflow.computation\n",
"def tff_return_one_noise(i):\n",
- " g=tf.random.Generator.from_seed(i)\n",
+ " g = tf.random.Generator.from_seed(i)\n",
+ "\n",
" @tf.function\n",
" def tf_return_one_noise():\n",
" return g.normal([])\n",
+ "\n",
" return tf_return_one_noise()\n",
"\n",
- "@tff.federated_computation\n",
+ "\n",
+ "@federated_language.federated_computation\n",
"def return_two_noise():\n",
" return (tff_return_one_noise(1), tff_return_one_noise(2))\n",
"\n",
+ "\n",
"n1, n2 = return_two_noise()\n",
"assert n1 != n2\n",
"print(n1, n2)"
@@ -509,20 +536,24 @@
"source": [
"@tff.tensorflow.computation\n",
"def tff_return_one_noise(i):\n",
- " g=tf.random.Generator.from_seed(i)\n",
+ " g = tf.random.Generator.from_seed(i)\n",
" weights = [\n",
- " tf.ones([2, 2], dtype=tf.float32),\n",
- " tf.constant([2], dtype=tf.float32)\n",
- " ]\n",
+ " tf.ones([2, 2], dtype=tf.float32),\n",
+ " tf.constant([2], dtype=tf.float32),\n",
+ " ]\n",
+ "\n",
" @tf.function\n",
" def tf_return_one_noise():\n",
" return tf.nest.map_structure(lambda x: g.normal(tf.shape(x)), weights)\n",
+ "\n",
" return tf_return_one_noise()\n",
"\n",
- "@tff.federated_computation\n",
+ "\n",
+ "@federated_language.federated_computation\n",
"def return_two_noise():\n",
" return (tff_return_one_noise(1), tff_return_one_noise(2))\n",
"\n",
+ "\n",
"n1, n2 = return_two_noise()\n",
"assert n1[1] != n2[1]\n",
"print('n1', n1)\n",
@@ -559,7 +590,8 @@
" # tf.timestamp returns microseconds as decimal places, thus scaling by 1e6.\n",
" return tf.cast(tf.timestamp() * 1e6, tf.int64)\n",
"\n",
- "class RandomSeedGenerator():\n",
+ "\n",
+ "class RandomSeedGenerator:\n",
"\n",
" def initialize(self, seed=None):\n",
" if seed is None:\n",
@@ -573,8 +605,10 @@
" def structure_next(self, state, nest_structure):\n",
" \"Returns seed in nested structure and the next state seed.\"\n",
" flat_structure = tf.nest.flatten(nest_structure)\n",
- " flat_seeds = [state + tf.constant([0, i], tf.int64) for\n",
- " i in range(len(flat_structure))]\n",
+ " flat_seeds = [\n",
+ " state + tf.constant([0, i], tf.int64)\n",
+ " for i in range(len(flat_structure))\n",
+ " ]\n",
" nest_seeds = tf.nest.pack_sequence_as(nest_structure, flat_seeds)\n",
" return nest_seeds, flat_seeds[-1] + tf.constant([0, 1], tf.int64)"
]
@@ -610,31 +644,39 @@
"source": [
"@tff.tensorflow.computation\n",
"def tff_return_one_noise(seed_state):\n",
- " g=RandomSeedGenerator()\n",
+ " g = RandomSeedGenerator()\n",
" weights = [\n",
- " tf.ones([2, 2], dtype=tf.float32),\n",
- " tf.constant([2], dtype=tf.float32)\n",
- " ]\n",
+ " tf.ones([2, 2], dtype=tf.float32),\n",
+ " tf.constant([2], dtype=tf.float32),\n",
+ " ]\n",
+ "\n",
" @tf.function\n",
" def tf_return_one_noise():\n",
" nest_seeds, updated_state = g.structure_next(seed_state, weights)\n",
- " nest_noise = tf.nest.map_structure(lambda x,s: tf.random.stateless_normal(\n",
- " shape=tf.shape(x), seed=s), weights, nest_seeds)\n",
+ " nest_noise = tf.nest.map_structure(\n",
+ " lambda x, s: tf.random.stateless_normal(shape=tf.shape(x), seed=s),\n",
+ " weights,\n",
+ " nest_seeds,\n",
+ " )\n",
" return nest_noise, updated_state\n",
+ "\n",
" return tf_return_one_noise()\n",
"\n",
+ "\n",
"@tff.tensorflow.computation\n",
"def tff_init_state():\n",
- " g=RandomSeedGenerator()\n",
+ " g = RandomSeedGenerator()\n",
" return g.initialize()\n",
"\n",
- "@tff.federated_computation\n",
+ "\n",
+ "@federated_language.federated_computation\n",
"def return_two_noise():\n",
" seed_state = tff_init_state()\n",
" n1, seed_state = tff_return_one_noise(seed_state)\n",
" n2, seed_state = tff_return_one_noise(seed_state)\n",
" return (n1, n2)\n",
"\n",
+ "\n",
"n1, n2 = return_two_noise()\n",
"assert n1[1] != n2[1]\n",
"print('n1', n1)\n",
diff --git a/docs/tutorials/simulations_with_accelerators.ipynb b/docs/tutorials/simulations_with_accelerators.ipynb
index 2a4a9ac6e..1b09d1d84 100644
--- a/docs/tutorials/simulations_with_accelerators.ipynb
+++ b/docs/tutorials/simulations_with_accelerators.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -83,7 +83,8 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
+ "!pip install --quite --upgrade federated_language\n",
"!pip install --quiet --upgrade tensorflow-federated\n",
"!pip install -U tensorboard_plugin_profile"
]
@@ -110,6 +111,7 @@
"import collections\n",
"import time\n",
"\n",
+ "import federated_language\n",
"import numpy as np\n",
"import tensorflow as tf\n",
"import tensorflow_federated as tff"
@@ -152,9 +154,12 @@
" raise ValueError('Cannot detect physical GPU device in TF')\n",
"# TODO: b/277213652 - Remove this call, as it doesn't work with C++ executor\n",
"tf.config.set_logical_device_configuration(\n",
- " gpu_devices[0], \n",
- " [tf.config.LogicalDeviceConfiguration(memory_limit=1024),\n",
- " tf.config.LogicalDeviceConfiguration(memory_limit=1024)])\n",
+ " gpu_devices[0],\n",
+ " [\n",
+ " tf.config.LogicalDeviceConfiguration(memory_limit=1024),\n",
+ " tf.config.LogicalDeviceConfiguration(memory_limit=1024),\n",
+ " ],\n",
+ ")\n",
"tf.config.list_logical_devices()"
]
},
@@ -190,10 +195,11 @@
}
],
"source": [
- "@tff.federated_computation\n",
+ "@federated_language.federated_computation\n",
"def hello_world():\n",
" return 'Hello, World!'\n",
"\n",
+ "\n",
"hello_world()"
]
},
@@ -216,7 +222,9 @@
},
"outputs": [],
"source": [
- "emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(only_digits=True)"
+ "emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(\n",
+ " only_digits=True\n",
+ ")"
]
},
{
@@ -236,24 +244,32 @@
},
"outputs": [],
"source": [
- "def preprocess_emnist_dataset(client_epochs_per_round, batch_size, test_batch_size):\n",
+ "def preprocess_emnist_dataset(\n",
+ " client_epochs_per_round, batch_size, test_batch_size\n",
+ "):\n",
"\n",
" def element_fn(element):\n",
" return collections.OrderedDict(\n",
- " x=tf.expand_dims(element['pixels'], -1), y=element['label'])\n",
+ " x=tf.expand_dims(element['pixels'], -1), y=element['label']\n",
+ " )\n",
"\n",
" def preprocess_train_dataset(dataset):\n",
" # Use buffer_size same as the maximum client dataset size,\n",
" # 418 for Federated EMNIST\n",
- " return dataset.map(element_fn).shuffle(buffer_size=418).repeat(\n",
- " count=client_epochs_per_round).batch(batch_size, drop_remainder=False)\n",
+ " return (\n",
+ " dataset.map(element_fn)\n",
+ " .shuffle(buffer_size=418)\n",
+ " .repeat(count=client_epochs_per_round)\n",
+ " .batch(batch_size, drop_remainder=False)\n",
+ " )\n",
"\n",
" def preprocess_test_dataset(dataset):\n",
" return dataset.map(element_fn).batch(test_batch_size, drop_remainder=False)\n",
"\n",
" train_set = emnist_train.preprocess(preprocess_train_dataset)\n",
" test_set = preprocess_test_dataset(\n",
- " emnist_test.create_tf_dataset_from_all_clients())\n",
+ " emnist_test.create_tf_dataset_from_all_clients()\n",
+ " )\n",
" return train_set, test_set"
]
},
@@ -304,13 +320,13 @@
" \"\"\"A stack of basic blocks.\"\"\"\n",
" x = _basic_block(input_tensor, filters, strides=strides)\n",
" for _ in range(size - 1):\n",
- " x = _basic_block(x, filters, strides=1)\n",
+ " x = _basic_block(x, filters, strides=1)\n",
" return x\n",
"\n",
"\n",
"def create_cnn(num_blocks, conv_width_multiplier=1, num_classes=10):\n",
- " \"\"\"Create a VGG-like CNN model. \n",
- " \n",
+ " \"\"\"Create a VGG-like CNN model.\n",
+ "\n",
" The CNN has (6*num_blocks + 2) layers.\n",
" \"\"\"\n",
" input_shape = (28, 28, 1) # channels_last\n",
@@ -319,9 +335,15 @@
" x = tf.image.per_image_standardization(x)\n",
"\n",
" x = _conv_3x3(x, 16 * conv_width_multiplier, 1)\n",
- " x = _vgg_block(x, size=num_blocks, filters=16 * conv_width_multiplier, strides=1)\n",
- " x = _vgg_block(x, size=num_blocks, filters=32 * conv_width_multiplier, strides=2)\n",
- " x = _vgg_block(x, size=num_blocks, filters=64 * conv_width_multiplier, strides=2)\n",
+ " x = _vgg_block(\n",
+ " x, size=num_blocks, filters=16 * conv_width_multiplier, strides=1\n",
+ " )\n",
+ " x = _vgg_block(\n",
+ " x, size=num_blocks, filters=32 * conv_width_multiplier, strides=2\n",
+ " )\n",
+ " x = _vgg_block(\n",
+ " x, size=num_blocks, filters=64 * conv_width_multiplier, strides=2\n",
+ " )\n",
"\n",
" x = tf.keras.layers.GlobalAveragePooling2D()(x)\n",
" x = tf.keras.layers.Dense(num_classes)(x)\n",
@@ -329,9 +351,9 @@
" model = tf.keras.models.Model(\n",
" img_input,\n",
" x,\n",
- " name='cnn-{}-{}'.format(6 * num_blocks + 2, conv_width_multiplier))\n",
- " return model\n",
- " "
+ " name='cnn-{}-{}'.format(6 * num_blocks + 2, conv_width_multiplier),\n",
+ " )\n",
+ " return model"
]
},
{
@@ -359,37 +381,42 @@
" return metric.result()\n",
"\n",
"\n",
- "def run_federated_training(client_epochs_per_round, \n",
- " train_batch_size, \n",
- " test_batch_size, \n",
- " cnn_num_blocks, \n",
- " conv_width_multiplier,\n",
- " server_learning_rate, \n",
- " client_learning_rate, \n",
- " total_rounds, \n",
- " clients_per_round, \n",
- " rounds_per_eval,\n",
- " logdir='logdir'):\n",
+ "def run_federated_training(\n",
+ " client_epochs_per_round,\n",
+ " train_batch_size,\n",
+ " test_batch_size,\n",
+ " cnn_num_blocks,\n",
+ " conv_width_multiplier,\n",
+ " server_learning_rate,\n",
+ " client_learning_rate,\n",
+ " total_rounds,\n",
+ " clients_per_round,\n",
+ " rounds_per_eval,\n",
+ " logdir='logdir',\n",
+ "):\n",
"\n",
" train_data, test_data = preprocess_emnist_dataset(\n",
- " client_epochs_per_round, train_batch_size, test_batch_size)\n",
+ " client_epochs_per_round, train_batch_size, test_batch_size\n",
+ " )\n",
" data_spec = test_data.element_spec\n",
"\n",
" def _model_fn():\n",
" keras_model = create_cnn(cnn_num_blocks, conv_width_multiplier)\n",
" loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
" return tff.learning.models.from_keras_model(\n",
- " keras_model, input_spec=data_spec, loss=loss)\n",
- " \n",
+ " keras_model, input_spec=data_spec, loss=loss\n",
+ " )\n",
+ "\n",
" server_optimizer = tff.learning.optimizers.build_sgdm(server_learning_rate)\n",
" client_optimizer = tff.learning.optimizers.build_sgdm(client_learning_rate)\n",
- " \n",
+ "\n",
" learning_process = tff.learning.algorithms.build_weighted_fed_avg(\n",
- " model_fn=_model_fn, \n",
- " server_optimizer_fn=server_optimizer, \n",
- " client_optimizer_fn=client_optimizer, \n",
- " use_experimental_simulation_loop=True)\n",
- " \n",
+ " model_fn=_model_fn,\n",
+ " server_optimizer_fn=server_optimizer,\n",
+ " client_optimizer_fn=client_optimizer,\n",
+ " use_experimental_simulation_loop=True,\n",
+ " )\n",
+ "\n",
" metric = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')\n",
" eval_model = create_cnn(cnn_num_blocks, conv_width_multiplier)\n",
" logging.info(eval_model.summary())\n",
@@ -398,29 +425,28 @@
" start_time = time.time()\n",
" for round_num in range(total_rounds):\n",
" sampled_clients = np.random.choice(\n",
- " train_data.client_ids,\n",
- " size=clients_per_round,\n",
- " replace=False)\n",
+ " train_data.client_ids, size=clients_per_round, replace=False\n",
+ " )\n",
" sampled_train_data = [\n",
" train_data.create_tf_dataset_for_client(client)\n",
" for client in sampled_clients\n",
" ]\n",
- " if round_num == total_rounds-1:\n",
+ " if round_num == total_rounds - 1:\n",
" with tf.profiler.experimental.Profile(logdir):\n",
- " result = learning_process.next(\n",
- " server_state, sampled_train_data)\n",
+ " result = learning_process.next(server_state, sampled_train_data)\n",
" else:\n",
- " result = learning_process.next(\n",
- " server_state, sampled_train_data)\n",
+ " result = learning_process.next(server_state, sampled_train_data)\n",
" server_state = result.state\n",
" train_metrics = result.metrics['client_work']['train']\n",
- " print(f'Round {round_num} training loss: {train_metrics[\"loss\"]}, '\n",
- " f'time: {(time.time()-start_time)/(round_num+1.)} secs')\n",
- " if round_num % rounds_per_eval == 0 or round_num == total_rounds-1:\n",
+ " print(\n",
+ " f'Round {round_num} training loss: {train_metrics[\"loss\"]}, '\n",
+ " f'time: {(time.time()-start_time)/(round_num+1.)} secs'\n",
+ " )\n",
+ " if round_num % rounds_per_eval == 0 or round_num == total_rounds - 1:\n",
" model_weights = learning_process.get_model_weights(server_state)\n",
" model_weights.assign_weights_to(eval_model)\n",
" accuracy = keras_evaluate(eval_model, test_data, metric)\n",
- " print(f'Round {round_num} validation accuracy: {accuracy * 100.0}')\n"
+ " print(f'Round {round_num} validation accuracy: {accuracy * 100.0}')"
]
},
{
@@ -531,17 +557,17 @@
],
"source": [
"run_federated_training(\n",
- " client_epochs_per_round=1, \n",
- " train_batch_size=16, \n",
- " test_batch_size=128, \n",
- " cnn_num_blocks=2, \n",
+ " client_epochs_per_round=1,\n",
+ " train_batch_size=16,\n",
+ " test_batch_size=128,\n",
+ " cnn_num_blocks=2,\n",
" conv_width_multiplier=4,\n",
- " server_learning_rate=1.0, \n",
+ " server_learning_rate=1.0,\n",
" client_learning_rate=0.01,\n",
" total_rounds=10,\n",
- " clients_per_round=16, \n",
+ " clients_per_round=16,\n",
" rounds_per_eval=2,\n",
- " )\n"
+ ")"
]
},
{
@@ -552,7 +578,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"%tensorboard --logdir=logdir --port=0"
]
},
@@ -704,17 +730,17 @@
],
"source": [
"run_federated_training(\n",
- " client_epochs_per_round=1, \n",
- " train_batch_size=16, \n",
- " test_batch_size=128, \n",
- " cnn_num_blocks=4, \n",
+ " client_epochs_per_round=1,\n",
+ " train_batch_size=16,\n",
+ " test_batch_size=128,\n",
+ " cnn_num_blocks=4,\n",
" conv_width_multiplier=4,\n",
- " server_learning_rate=1.0, \n",
+ " server_learning_rate=1.0,\n",
" client_learning_rate=0.01,\n",
" total_rounds=5,\n",
- " clients_per_round=16, \n",
+ " clients_per_round=16,\n",
" rounds_per_eval=2,\n",
- " )"
+ ")"
]
},
{
@@ -868,20 +894,21 @@
"source": [
"# Control concurrency by `max_concurrent_computation_calls`.\n",
"tff.backends.native.set_sync_local_cpp_execution_context(\n",
- " max_concurrent_computation_calls=16/2)\n",
+ " max_concurrent_computation_calls=16 / 2\n",
+ ")\n",
"\n",
"run_federated_training(\n",
- " client_epochs_per_round=1, \n",
- " train_batch_size=16, \n",
- " test_batch_size=128, \n",
- " cnn_num_blocks=4, \n",
+ " client_epochs_per_round=1,\n",
+ " train_batch_size=16,\n",
+ " test_batch_size=128,\n",
+ " cnn_num_blocks=4,\n",
" conv_width_multiplier=4,\n",
- " server_learning_rate=1.0, \n",
+ " server_learning_rate=1.0,\n",
" client_learning_rate=0.01,\n",
" total_rounds=5,\n",
- " clients_per_round=16, \n",
+ " clients_per_round=16,\n",
" rounds_per_eval=2,\n",
- " )"
+ ")"
]
},
{
@@ -1032,24 +1059,24 @@
}
],
"source": [
- "# Mixed precision training. \n",
+ "# Mixed precision training.\n",
"tff.backends.native.set_sync_local_cpp_execution_context()\n",
"policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')\n",
"tf.keras.mixed_precision.experimental.set_policy(policy)\n",
"\n",
"run_federated_training(\n",
- " client_epochs_per_round=1, \n",
- " train_batch_size=16, \n",
- " test_batch_size=128, \n",
- " cnn_num_blocks=4, \n",
+ " client_epochs_per_round=1,\n",
+ " train_batch_size=16,\n",
+ " test_batch_size=128,\n",
+ " cnn_num_blocks=4,\n",
" conv_width_multiplier=4,\n",
- " server_learning_rate=1.0, \n",
+ " server_learning_rate=1.0,\n",
" client_learning_rate=0.01,\n",
" total_rounds=5,\n",
- " clients_per_round=16, \n",
+ " clients_per_round=16,\n",
" rounds_per_eval=2,\n",
- " logdir='mixed'\n",
- " )"
+ " logdir='mixed',\n",
+ ")"
]
},
{
@@ -1060,7 +1087,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"%tensorboard --logdir=mixed --port=0"
]
}
diff --git a/docs/tutorials/tuning_recommended_aggregators.ipynb b/docs/tutorials/tuning_recommended_aggregators.ipynb
index eeb84065f..9187ff38d 100644
--- a/docs/tutorials/tuning_recommended_aggregators.ipynb
+++ b/docs/tutorials/tuning_recommended_aggregators.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -95,7 +95,7 @@
},
"outputs": [],
"source": [
- "#@test {\"skip\": true}\n",
+ "# @test {\"skip\": true}\n",
"!pip install --quiet --upgrade tensorflow-federated"
]
},
@@ -120,7 +120,8 @@
"source": [
"import math\n",
"import tensorflow_federated as tff\n",
- "tff.federated_computation(lambda: 'Hello, World!')()"
+ "\n",
+ "tff.tensorflow_computation(lambda: 'Hello, World!')()"
]
},
{
@@ -198,7 +199,8 @@
"outputs": [],
"source": [
"median_estimate = tff.aggregators.PrivateQuantileEstimationProcess.no_noise(\n",
- " initial_estimate=1.0, target_quantile=0.5, learning_rate=0.2)"
+ " initial_estimate=1.0, target_quantile=0.5, learning_rate=0.2\n",
+ ")"
]
},
{
@@ -260,10 +262,11 @@
" target_quantile=0.98,\n",
" learning_rate=math.log(10),\n",
" multiplier=2.0,\n",
- " increment=1.0)\n",
+ " increment=1.0,\n",
+ ")\n",
"zeroing_mean = tff.aggregators.zeroing_factory(\n",
- " zeroing_norm=zeroing_norm,\n",
- " inner_agg_factory=tff.aggregators.MeanFactory())\n",
+ " zeroing_norm=zeroing_norm, inner_agg_factory=tff.aggregators.MeanFactory()\n",
+ ")\n",
"\n",
"# Equivalent to:\n",
"# zeroing_mean = tff.learning.robust_aggregator(clipping=False)"
@@ -298,12 +301,11 @@
"outputs": [],
"source": [
"clipping_norm = tff.aggregators.PrivateQuantileEstimationProcess.no_noise(\n",
- " initial_estimate=1.0,\n",
- " target_quantile=0.8,\n",
- " learning_rate=0.2)\n",
+ " initial_estimate=1.0, target_quantile=0.8, learning_rate=0.2\n",
+ ")\n",
"clipping_mean = tff.aggregators.clipping_factory(\n",
- " clipping_norm=clipping_norm,\n",
- " inner_agg_factory=tff.aggregators.MeanFactory())\n",
+ " clipping_norm=clipping_norm, inner_agg_factory=tff.aggregators.MeanFactory()\n",
+ ")\n",
"\n",
"# Equivalent to:\n",
"# clipping_mean = tff.learning.robust_aggregator(zeroing=False)"
@@ -338,7 +340,8 @@
"outputs": [],
"source": [
"dp_mean = tff.aggregators.DifferentiallyPrivateFactory.gaussian_adaptive(\n",
- " noise_multiplier=0.1, clients_per_round=100)\n",
+ " noise_multiplier=0.1, clients_per_round=100\n",
+ ")\n",
"\n",
"# Equivalent to:\n",
"# dp_mean = tff.learning.dp_aggregator(\n",
@@ -383,7 +386,9 @@
"source": [
"compressed_mean = tff.aggregators.MeanFactory(\n",
" tff.aggregators.EncodedSumFactory.quantize_above_threshold(\n",
- " quantization_bits=8, threshold=20000))\n",
+ " quantization_bits=8, threshold=20000\n",
+ " )\n",
+ ")\n",
"\n",
"# Equivalent to:\n",
"# compressed_mean = tff.learning.compression_aggregator(zeroing=False, clipping=False)"
@@ -456,9 +461,11 @@
" initial_estimate=50.0,\n",
" target_quantile=0.95,\n",
" learning_rate=1.0,\n",
- " multiplier=2.0)\n",
+ " multiplier=2.0,\n",
+ ")\n",
"secure_mean = tff.aggregators.MeanFactory(\n",
- " tff.aggregators.SecureSumFactory(secagg_bound))\n",
+ " tff.aggregators.SecureSumFactory(secagg_bound)\n",
+ ")\n",
"\n",
"# Equivalent to:\n",
"# secure_mean = tff.learning.secure_aggregator(zeroing=Fasle, clipping=False)"
diff --git a/docs/tutorials/working_with_client_data.ipynb b/docs/tutorials/working_with_client_data.ipynb
index 16a2ffffa..0c5ea2af7 100644
--- a/docs/tutorials/working_with_client_data.ipynb
+++ b/docs/tutorials/working_with_client_data.ipynb
@@ -18,7 +18,7 @@
},
"outputs": [],
"source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
@@ -105,8 +105,8 @@
},
"outputs": [],
"source": [
- "#@title Set up open-source environment\n",
- "#@test {\"skip\": true}\n",
+ "# @title Set up open-source environment\n",
+ "# @test {\"skip\": true}\n",
"\n",
"# tensorflow_federated_nightly also bring in tf_nightly, which\n",
"# can causes a duplicate tensorboard install, leading to errors.\n",
@@ -124,7 +124,7 @@
},
"outputs": [],
"source": [
- "#@title Import packages\n",
+ "# @title Import packages\n",
"import collections\n",
"import time\n",
"\n",
@@ -182,8 +182,7 @@
],
"source": [
"first_client_id = client_data.client_ids[0]\n",
- "first_client_dataset = client_data.create_tf_dataset_for_client(\n",
- " first_client_id)\n",
+ "first_client_dataset = client_data.create_tf_dataset_for_client(first_client_id)\n",
"print(first_client_dataset.element_spec)\n",
"# This information is also available as a `ClientData` property:\n",
"assert client_data.element_type_structure == first_client_dataset.element_spec"
@@ -224,15 +223,19 @@
" y=tf.cast(tf.reshape(input['label'], shape=(-1, 1)), tf.int64),\n",
" )\n",
"\n",
- " return dataset.batch(5).map(\n",
- " map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE).take(5)\n",
+ " return (\n",
+ " dataset.batch(5)\n",
+ " .map(map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
+ " .take(5)\n",
+ " )\n",
"\n",
"\n",
"preprocessed_client_data = client_data.preprocess(preprocess_dataset)\n",
"\n",
"# Notice that we have both reshaped and renamed the elements of the ordered dict.\n",
"first_client_dataset = preprocessed_client_data.create_tf_dataset_for_client(\n",
- " first_client_id)\n",
+ " first_client_id\n",
+ ")\n",
"print(first_client_dataset.element_spec)"
]
},
@@ -267,11 +270,13 @@
" preprocessed = preprocess_dataset(dataset)\n",
" return preprocessed.shuffle(buffer_size=5)\n",
"\n",
+ "\n",
"preprocessed_and_shuffled = client_data.preprocess(preprocess_and_shuffle)\n",
"\n",
"# The type signature will remain the same, but the batches will be shuffled.\n",
"first_client_dataset = preprocessed_and_shuffled.create_tf_dataset_for_client(\n",
- " first_client_id)\n",
+ " first_client_id\n",
+ ")\n",
"print(first_client_dataset.element_spec)"
]
},
@@ -463,8 +468,11 @@
},
"outputs": [],
"source": [
- "trainer_accepting_ids = tff.simulation.compose_dataset_computation_with_iterative_process(\n",
- " preprocessed_and_shuffled.dataset_computation, trainer)"
+ "trainer_accepting_ids = (\n",
+ " tff.simulation.compose_dataset_computation_with_iterative_process(\n",
+ " preprocessed_and_shuffled.dataset_computation, trainer\n",
+ " )\n",
+ ")"
]
},
{
@@ -531,7 +539,7 @@
")\n",
"\n",
"\n",
- "@tff.federated_computation(\n",
+ "@federated_language(\n",
" trainer.next.type_signature.parameter[0], selected_clients_type\n",
")\n",
"def new_next(server_state, selected_clients):\n",