<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2020 The TensorFlow Authors.
Step1:  フェデレーテッドラーニングリサーチの TFF
Step2:  TFF が動作していることを確認します。
Step4:  入力データを準備する
このセクションでは、TFF に含まれる EMNIST データセットを読み込んで事前処理します。EMNIST データセットの詳細は、画像分類のフェデレーテッドラーニングチュートリアルをご覧ください。
Step6:  モデルを定義する
ここでは、元の FedAvg CNN に基づいて Keras モデルを定義し、それを tff.learning.Model インスタンスにラッピングして TFF が消費できるようにします。
モデルのみを直接生成する代わりに、モデルを生成する関数が必要となることに注意してください。また、その関数は構築済みのモデルをキャプチャするだけでなく、呼び出されるコンテキストで作成する必要があります。これは、TFF がデバイスで利用されるように設計されており、リソースが作られるタイミングを制御することで、キャプチャしてパッケージ化できる必要があるためです。
Step7:  モデルのトレーニングとトレーニングメトリックの出力
フェデレーテッドアベレージングアルゴリズムを作成し、定義済みのモデルを EMNIST データセットでトレーニングする準備が整いました。
まず、tff.learning.build_federated_averaging_process API を使用して、フェデレーテッドアベレージングアルゴリズムを構築する必要があります。
Step11:  では、フェデレーテッドアベレージングアルゴリズムを実行しましょう。TFF の観点からフェデレーテッドアベレージングアルゴリズムを実行するには、次のようになります。
アルゴリズムを初期化し、サーバーの初期状態を取得します。サーバーの状態には、アルゴリズムを実行するために必要な情報が含まれます。TFF は関数型であるため、この状態には、アルゴリズムが使用するオプティマイザの状態（慣性項）だけでなく、モデルパラメータ自体も含まれることを思い出してください。これらは引数として渡され、TFF 計算の結果として返されます。
ラウンドごとにアルゴリズムを実行します。各ラウンドでは、新しいサーバーの状態が、データでモデルをトレーニングしている各クライアントの結果として返されます。通常、1 つのラウンドでは次のことが発生します。
サーバーはすべての参加クライアントにモデルをブロードキャストします。
各クライアントは、モデルとそのデータに基づいて作業を実施します。
サーバーはすべてのモデルを集約し、新しいモデルを含むサーバーの状態を生成します。
詳細については、カスタムフェデレーテッドアルゴリズム、パート 2
Step12:  上記に示されるルートログディレクトリで TensorBoard を起動すると、トレーニングメトリックが表示されます。データの読み込みには数秒かかることがあります。Loss と Accuracy を除き、ブロードキャストされ集約されたデータの量も出力されます。ブロードキャストされたデータは、各クライアントにサーバーがプッシュしたテンソルで、集約データとは各クライアントがサーバーに返すテンソルを指します。
Step15:  カスタムブロードキャストと集約関数を構築する
tensor_encoding API を使用して、ブロードキャストされたデータと集約データに対して非可逆圧縮アルゴリズムを使用する関数を実装しましょう。
まず、2 つの関数を定義します。
broadcast_encoder_fn
Step16:  TFF は、エンコーダ関数を tff.learning.build_federated_averaging_process API が消費できる形式に変換する API を提供しています。tff.learning.framework.build_encoded_broadcast_from_model と tff.learning.framework.build_encoded_mean_from_model を使用することで、tff.learning.build_federated_averaging_process の broadcast_process と aggregation_process 引数に渡して、非可逆圧縮アルゴリズムでフェデレーテッドアベレージングアルゴリズムを作成するための関数を 2 つ作成することができます。
Step17:  もう一度モデルをトレーニングする
では、新しいフェデレーテッドアベレージングアルゴリズムを実行しましょう。
Step18:  もう一度 TensorBoard を起動して、2 つの実行のトレーニングメトリックを比較します。
Tensorboard を見てわかるように、broadcasted_bits と aggregated_bits 図 の orginial と compression の曲線に大きな減少を確認できます。loss と sparse_categorical_accuracy 図では、この 2 つの曲線は非常に似通っていました。
最後に、元のフェデレーテッドアベレージングアルゴリズムに似たパフォーマンスを達成できる圧縮アルゴリズムを実装しながら、通信コストを大幅に削減することができました。
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2020 The TensorFlow Authors.
End of explanation
#@test {"skip": true}
!pip install --quiet --upgrade tensorflow_federated
!pip install --quiet --upgrade tensorflow-model-optimization
%load_ext tensorboard
import functools
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
from tensorflow_model_optimization.python.core.internal import tensor_encoding as te
Explanation: フェデレーテッドラーニングリサーチの TFF: モデルと更新圧縮
注意: この Colab は <a>最新リリースバージョン</a>の <code>tensorflow_federated</code> pip パッケージでの動作が確認されていますが、Tensorflow Federated プロジェクトは現在もプレリリース開発の段階にあるため、master では動作しない可能性があります。
このチュートリアルでは、EMNIST データセットを使用しながら、tff.learning.build_federated_averaging_process API と tensor_encoding API を使用するフェデレーテッドアベレージングアルゴリズムにおける通信コストを削減するために非可逆圧縮アルゴリズムを有効化する方法を実演します。フェデレーテッドアベレージングアルゴリズムの詳細については、論文「Communication-Efficient Learning of Deep Networks from Decentralized Data」をご覧ください。
始める前に
始める前に、次のコードを実行し、環境が正しくセットアップされていることを確認してください。挨拶文が表示されない場合は、インストールガイドで手順を確認してください。
End of explanation
@tff.federated_computation
def hello_world():
  return 'Hello, World!'
hello_world()
Explanation: TFF が動作していることを確認します。
End of explanation
# This value only applies to EMNIST dataset, consider choosing appropriate
# values if switching to other datasets.
MAX_CLIENT_DATASET_SIZE = 418
CLIENT_EPOCHS_PER_ROUND = 1
CLIENT_BATCH_SIZE = 20
TEST_BATCH_SIZE = 500
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(
    only_digits=True)
def reshape_emnist_element(element):
  return (tf.expand_dims(element['pixels'], axis=-1), element['label'])
def preprocess_train_dataset(dataset):
  Preprocessing function for the EMNIST training dataset.
  return (dataset
          # Shuffle according to the largest client dataset
          .shuffle(buffer_size=MAX_CLIENT_DATASET_SIZE)
          # Repeat to do multiple local epochs
          .repeat(CLIENT_EPOCHS_PER_ROUND)
          # Batch to a fixed client batch size
          .batch(CLIENT_BATCH_SIZE, drop_remainder=False)
          # Preprocessing step
          .map(reshape_emnist_element))
emnist_train = emnist_train.preprocess(preprocess_train_dataset)
Explanation: 入力データを準備する
このセクションでは、TFF に含まれる EMNIST データセットを読み込んで事前処理します。EMNIST データセットの詳細は、画像分類のフェデレーテッドラーニングチュートリアルをご覧ください。
End of explanation
def create_original_fedavg_cnn_model(only_digits=True):
  The CNN model used in https://arxiv.org/abs/1602.05629.
  data_format = 'channels_last'
  max_pool = functools.partial(
      tf.keras.layers.MaxPooling2D,
      pool_size=(2, 2),
      padding='same',
      data_format=data_format)
  conv2d = functools.partial(
      tf.keras.layers.Conv2D,
      kernel_size=5,
      padding='same',
      data_format=data_format,
      activation=tf.nn.relu)
  model = tf.keras.models.Sequential([
      tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
      conv2d(filters=32),
      max_pool(),
      conv2d(filters=64),
      max_pool(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(512, activation=tf.nn.relu),
      tf.keras.layers.Dense(10 if only_digits else 62),
      tf.keras.layers.Softmax(),
  ])
  return model
# Gets the type information of the input data. TFF is a strongly typed
# functional programming framework, and needs type information about inputs to 
# the model.
input_spec = emnist_train.create_tf_dataset_for_client(
    emnist_train.client_ids[0]).element_spec
def tff_model_fn():
  keras_model = create_original_fedavg_cnn_model()
  return tff.learning.from_keras_model(
      keras_model=keras_model,
      input_spec=input_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
Explanation: モデルを定義する
ここでは、元の FedAvg CNN に基づいて Keras モデルを定義し、それを tff.learning.Model インスタンスにラッピングして TFF が消費できるようにします。
モデルのみを直接生成する代わりに、モデルを生成する関数が必要となることに注意してください。また、その関数は構築済みのモデルをキャプチャするだけでなく、呼び出されるコンテキストで作成する必要があります。これは、TFF がデバイスで利用されるように設計されており、リソースが作られるタイミングを制御することで、キャプチャしてパッケージ化できる必要があるためです。
End of explanation
federated_averaging = tff.learning.build_federated_averaging_process(
    model_fn=tff_model_fn,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))
Explanation: モデルのトレーニングとトレーニングメトリックの出力
フェデレーテッドアベレージングアルゴリズムを作成し、定義済みのモデルを EMNIST データセットでトレーニングする準備が整いました。
まず、tff.learning.build_federated_averaging_process API を使用して、フェデレーテッドアベレージングアルゴリズムを構築する必要があります。
End of explanation
#@title Load utility functions
def format_size(size):
  A helper function for creating a human-readable size.
  size = float(size)
  for unit in ['B','KiB','MiB','GiB']:
    if size < 1024.0:
      return "{size:3.2f}{unit}".format(size=size, unit=unit)
    size /= 1024.0
  return "{size:.2f}{unit}".format(size=size, unit='TiB')
def set_sizing_environment():
  Creates an environment that contains sizing information.
  # Creates a sizing executor factory to output communication cost
  # after the training finishes. Note that sizing executor only provides an
  # estimate (not exact) of communication cost, and doesn't capture cases like
  # compression of over-the-wire representations. However, it's perfect for
  # demonstrating the effect of compression in this tutorial.
  sizing_factory = tff.framework.sizing_executor_factory()
  # TFF has a modular runtime you can configure yourself for various
  # environments and purposes, and this example just shows how to configure one
  # part of it to report the size of things.
  context = tff.framework.ExecutionContext(executor_fn=sizing_factory)
  tff.framework.set_default_context(context)
  return sizing_factory
def train(federated_averaging_process, num_rounds, num_clients_per_round, summary_writer):
  Trains the federated averaging process and output metrics.
  # Create a environment to get communication cost.
  environment = set_sizing_environment()
  # Initialize the Federated Averaging algorithm to get the initial server state.
  state = federated_averaging_process.initialize()
  with summary_writer.as_default():
    for round_num in range(num_rounds):
      # Sample the clients parcitipated in this round.
      sampled_clients = np.random.choice(
          emnist_train.client_ids,
          size=num_clients_per_round,
          replace=False)
      # Create a list of `tf.Dataset` instances from the data of sampled clients.
      sampled_train_data = [
          emnist_train.create_tf_dataset_for_client(client)
          for client in sampled_clients
      ]
      # Round one round of the algorithm based on the server state and client data
      # and output the new state and metrics.
      state, metrics = federated_averaging_process.next(state, sampled_train_data)
      # For more about size_info, please see https://www.tensorflow.org/federated/api_docs/python/tff/framework/SizeInfo
      size_info = environment.get_size_info()
      broadcasted_bits = size_info.broadcast_bits[-1]
      aggregated_bits = size_info.aggregate_bits[-1]
      print('round {:2d}, metrics={}, broadcasted_bits={}, aggregated_bits={}'.format(round_num, metrics, format_size(broadcasted_bits), format_size(aggregated_bits)))
      # Add metrics to Tensorboard.
      for name, value in metrics['train']._asdict().items():
          tf.summary.scalar(name, value, step=round_num)
      # Add broadcasted and aggregated data size to Tensorboard.
      tf.summary.scalar('cumulative_broadcasted_bits', broadcasted_bits, step=round_num)
      tf.summary.scalar('cumulative_aggregated_bits', aggregated_bits, step=round_num)
      summary_writer.flush()
# Clean the log directory to avoid conflicts.
!rm -R /tmp/logs/scalars/*
# Set up the log directory and writer for Tensorboard.
logdir = "/tmp/logs/scalars/original/"
summary_writer = tf.summary.create_file_writer(logdir)
train(federated_averaging_process=federated_averaging, num_rounds=10,
      num_clients_per_round=10, summary_writer=summary_writer)
Explanation: では、フェデレーテッドアベレージングアルゴリズムを実行しましょう。TFF の観点からフェデレーテッドアベレージングアルゴリズムを実行するには、次のようになります。
アルゴリズムを初期化し、サーバーの初期状態を取得します。サーバーの状態には、アルゴリズムを実行するために必要な情報が含まれます。TFF は関数型であるため、この状態には、アルゴリズムが使用するオプティマイザの状態（慣性項）だけでなく、モデルパラメータ自体も含まれることを思い出してください。これらは引数として渡され、TFF 計算の結果として返されます。
ラウンドごとにアルゴリズムを実行します。各ラウンドでは、新しいサーバーの状態が、データでモデルをトレーニングしている各クライアントの結果として返されます。通常、1 つのラウンドでは次のことが発生します。
サーバーはすべての参加クライアントにモデルをブロードキャストします。
各クライアントは、モデルとそのデータに基づいて作業を実施します。
サーバーはすべてのモデルを集約し、新しいモデルを含むサーバーの状態を生成します。
詳細については、カスタムフェデレーテッドアルゴリズム、パート 2: フェデレーテッドアベレージングの実装チュートリアルをご覧ください。
トレーニングメトリックは、トレーニング後に表示できるように、TensorBoard ディレクトリに書き込まれます。
End of explanation
%tensorboard --logdir /tmp/logs/scalars/ --port=0
Explanation: 上記に示されるルートログディレクトリで TensorBoard を起動すると、トレーニングメトリックが表示されます。データの読み込みには数秒かかることがあります。Loss と Accuracy を除き、ブロードキャストされ集約されたデータの量も出力されます。ブロードキャストされたデータは、各クライアントにサーバーがプッシュしたテンソルで、集約データとは各クライアントがサーバーに返すテンソルを指します。
End of explanation
def broadcast_encoder_fn(value):
  Function for building encoded broadcast.
  spec = tf.TensorSpec(value.shape, value.dtype)
  if value.shape.num_elements() > 10000:
    return te.encoders.as_simple_encoder(
        te.encoders.uniform_quantization(bits=8), spec)
  else:
    return te.encoders.as_simple_encoder(te.encoders.identity(), spec)
def mean_encoder_fn(value):
  Function for building encoded mean.
  spec = tf.TensorSpec(value.shape, value.dtype)
  if value.shape.num_elements() > 10000:
    return te.encoders.as_gather_encoder(
        te.encoders.uniform_quantization(bits=8), spec)
  else:
    return te.encoders.as_gather_encoder(te.encoders.identity(), spec)
Explanation: カスタムブロードキャストと集約関数を構築する
tensor_encoding API を使用して、ブロードキャストされたデータと集約データに対して非可逆圧縮アルゴリズムを使用する関数を実装しましょう。
まず、2 つの関数を定義します。
broadcast_encoder_fn: サーバーのテンソルまたは変数をクライアント通信にエンコードする te.core.SimpleEncoder のインスタンスを作成します（ブロードキャストデータ）。
mean_encoder_fn: クライアントのテンソルまたは変数をサーバー通信にエンコードする te.core.GatherEncoder インスタンスを作成します（集約データ）。
一度にモデル全体に圧縮メソッドを適用しないことに十分に注意してください。モデルの各変数を圧縮するかどうか、またはどのように圧縮するかは、個別に決定します。これは一般的に、バイアスなどの小さな変数は不正確性により敏感であり、比較的小さいことから、潜在的な通信の節約量も比較的小さくなるためです。そのため、デフォルトでは小さな変数を圧縮しません。この例では、10000 個を超える要素を持つ変数ごとに 8 ビット（256 バケット）の均一量子化を適用し、ほかの変数にのみ ID を適用します。
End of explanation
encoded_broadcast_process = (
    tff.learning.framework.build_encoded_broadcast_process_from_model(
        tff_model_fn, broadcast_encoder_fn))
encoded_mean_process = (
    tff.learning.framework.build_encoded_mean_process_from_model(
    tff_model_fn, mean_encoder_fn))
federated_averaging_with_compression = tff.learning.build_federated_averaging_process(
    tff_model_fn,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0),
    broadcast_process=encoded_broadcast_process,
    aggregation_process=encoded_mean_process)
Explanation: TFF は、エンコーダ関数を tff.learning.build_federated_averaging_process API が消費できる形式に変換する API を提供しています。tff.learning.framework.build_encoded_broadcast_from_model と tff.learning.framework.build_encoded_mean_from_model を使用することで、tff.learning.build_federated_averaging_process の broadcast_process と aggregation_process 引数に渡して、非可逆圧縮アルゴリズムでフェデレーテッドアベレージングアルゴリズムを作成するための関数を 2 つ作成することができます。
End of explanation
logdir_for_compression = "/tmp/logs/scalars/compression/"
summary_writer_for_compression = tf.summary.create_file_writer(
    logdir_for_compression)
train(federated_averaging_process=federated_averaging_with_compression, 
      num_rounds=10,
      num_clients_per_round=10,
      summary_writer=summary_writer_for_compression)
Explanation: もう一度モデルをトレーニングする
では、新しいフェデレーテッドアベレージングアルゴリズムを実行しましょう。
End of explanation
%tensorboard --logdir /tmp/logs/scalars/ --port=0
Explanation: もう一度 TensorBoard を起動して、2 つの実行のトレーニングメトリックを比較します。
Tensorboard を見てわかるように、broadcasted_bits と aggregated_bits 図 の orginial と compression の曲線に大きな減少を確認できます。loss と sparse_categorical_accuracy 図では、この 2 つの曲線は非常に似通っていました。
最後に、元のフェデレーテッドアベレージングアルゴリズムに似たパフォーマンスを達成できる圧縮アルゴリズムを実装しながら、通信コストを大幅に削減することができました。
End of explanation
<END_TASK>