カスタムフェデレーテッドアルゴリズム、パート 1: フェデレーテッドコアの基礎

TensorFlow.orgで表示 Google Colab で実行 GitHub でソースを表示{

本チュートリアルは、フェデレーテッドコア (FC) を使用してTensorFlow Federated (TFF) でフェデレーテッドアルゴリズムのカスタム型を実装する方法を示す、2 部構成シリーズの第 1 部です。これはフェデレーテッドラーニング (FL) レイヤー実装の基礎として機能する、低レベルのインターフェースセットです。

この第 1 部はより概念的な内容です。TFF で使用している主要な概念およびプログラミングの抽象化をいくつか紹介し、それを温度センサーの分散配列を用いて非常に単純な例に使用する方法を示します。このシリーズの第 2 部では、ここで紹介する仕組みを使用してフェデレーテッドトレーニングと評価アルゴリズムの単純なバージョンを実装します。フォローアップとしては、tff.learningを使用したフェデレーテッドアベレージング実装の研究をお勧めします。

このシリーズを終えると、フェデレーテッドコア (FC) のアプリケーションが学習のみに限定されているわけではないことが認識できるはずです。提供しているプログラミングの抽象化は非常に一般的なので、例えば分散データに対する分析やその他カスタム型の計算の実装に使用することができます。

本チュートリアルは自己完結型で作成していますが、まず画像分類テキスト作成 に関するチュートリアルを読むことをお勧めします。TensorFlow Federated のフレームワークとフェデレーテッドラーニング API(tff.learning)に関する、より高レベルで丁寧な導入を提供しているので、ここで説明する概念を文脈に沿って理解するのに有用です。

使用目的

端的に言うと、フェデレーテッドコア (FC) は TensorFlow のコードと分散通信演算子を結合し、コンパクトなプログラムロジックの表現を可能にする開発環境です。これはフェデレーテッドアベレージングで使用されている、システム内のクライアントデバイスの集合に対する分散加算、平均、その他の型の分散集約の計算、およびそれらのデバイスのモデルやパラメータのブロードキャスト、などがそれにあたります。

tf.contrib.distributeについては既にご存知かも知れません。この時点で「このフレームワークはどのような点で違うのだろうか?」という自然な疑問が出てくるでしょう。結局のところは、どちらのフレームワークも TensorFlow の計算の分散を試みます。

1 つの考え方としては、tf.contrib.distributeの目標がユーザーが既存のモデルやトレーニングコードを最小限の変更で使用して分散トレーニングを可能にすることであるのに対し、TFF のフェデレーテッドコアの目標はシステムで使用する分散通信の特定のパターンを研究者や実践者が明示的に制御できるようにすることです。FC は、実装された分散トレーニング機能の具体的なセットの提供よりも、分散データフローアルゴリズムを表現するための柔軟で拡張可能な言語の提供に焦点を絞っています。

TFF の FC API の主要な対象者には、システム実装の詳細で行き詰まることなく新しいフェデレーテッドラーニングアルゴリズムを実験し、分散システム内でデータのフローのオーケストレーションに影響を与える設計の微妙な選択の結果を評価したいと考えている研究者や実践者があります。FC API が目指している抽象度は、システムに存在するデータとその変換方法など、研究出版物の中でフェデレーテッドラーニングアルゴリズムの仕組みの説明に使用される疑似コードにほぼ対応していますが、個々のポイントツーポイントのネットワークメッセージ交換のレベルにまで落ちることはありません。

TFF 全体はデータが分散しているシナリオを対象としており、例えばプライバシー上の理由など、データが分散した状態を維持する必要があるため、すべてのデータを中央の場所に収集することは実行不可能な選択肢である場合があります。これはすべてのデータをデータセンターの中央の場所に集められるシナリオと比較して、より高度で明示的な制御を必要とする機械学習アルゴリズムを実装する意義があると言えます。

始める前に

コードに手をつける前に、まず以下の「Hello World」の例を実行して、環境が正しく設定されていることを確認してください。動作しない場合は、インストールガイドをご覧ください。

pip install --quiet --upgrade tensorflow_federated
import collections

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

hello_world()
b'Hello, World!'

フェデレーテッドデータ

TFF の際立った特徴の 1 つは、フェデレーテッドデータに関する TensorFlow ベースの計算をコンパクトに表現できることです。本チュートリアルで使用するフェデレーテッドデータという用語は、分散システム内のデバイスのグループにまたがってホストされるデータアイテムの集まりを指します。例えば、モバイルデバイスで実行するアプリケーションはデータを収集し、中央の場所にはアップロードせずローカルに保存します。あるいは、分散センサーのアレイがその場所の温度測定値を収集して保存する場合などがあります。

上記の例のようなフェデレーテッドデータは、TFF では第一級オブジェクトとして扱います。つまり、それらは関数のパラメータおよび結果として表示され、型を持ちます。この概念を強化するために、フェデレーテッドデータセットをフェデレーテッド値、またはフェデレーテッド型の値と呼びます。

理解しておくべき重要な点は、すべてのデバイスにまたがるデータアイテムのコレクション全体(例えば、分散アレイ内の全センサーの温度測定値のコレクション全体)を単一のフェデレーテッド値としてモデル化するということです。

例として、クライアントデバイスのグループがホストするフェデレーテッド float 型を TFF で定義する方法を以下に示します。分散センサーのアレイにまたがってマテリアライズする温度測定値のコレクションは、このフェデレーテッドタイプの値としてモデル化されます。

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

より一般的には、TFF のフェデレーテッド型は、そのメンバ構成要素T型を指定して定義します。これには個々のデバイスに存在するデータのアイテムと、この型のフェデレーテッド値がホストされるデバイスのグループG(それに加えて後で説明するオプションの 3 つ目の情報)があります。フェデレーテッド値をホストするデバイスのグループGを、その値の配置と呼びます。したがって、tff.CLIENTSは配置の一例です。

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

以下に示すように、メンバ構成要素Tと配置Gを持つフェデレーテッド型は、コンパクトに{T}@Gと表すことができます。

str(federated_float_on_clients)
'{float32}@CLIENTS'

この簡潔な表記法である中括弧{}には、例えば温度センサーの測定値のようにメンバ構成要素(異なるデバイス上のデータのアイテム)が異なる場合があるので、クライアントがグループとしてT型のアイテムのマルチセットを共同でホストし、それらが一緒になってフェデレーテッド値を構成することを思い出させる役割があります。

重要な点として、フェデレーテッド値のメンバ構成要素は一般にプログラマには不透明だということがあります。つまり、フェデレーテッド値をシステム内のデバイスの識別子によってキー付けされる単純なdictであると考えるべきではありません。これらの値は、さまざまな種類の分散通信プロトコル(集約など)を抽象的に表現するフェデレーテッド演算子によってのみ、集合的に変換されるように意図されています。これが抽象的に見えたとしても心配は不要です。これについては、後ほど具体的な例を用いて説明します。

TFF のフェデレーテッド型には 2 つの種類があります。フェデレーテッド値のメンバ構成要素が(上記のように)異なる可能性があるものと、それらが全て等しいと分かっているものです。これは、tff.FederatedTypeコンストラクタの 3 番目のオプションであるall_equalパラメータによって制御されます(デフォルトは False)。

federated_float_on_clients.all_equal
False

T型のメンバ構成要素がすべて等しいと分かっている配置Gを持つフェデレーテッド型は、T@Gとしてコンパクトに表現できます(これは{T}@Gとは対照的で、メンバ構成要素のマルチセットが 1 つのアイテムで構成されているという事実を反映するために、中括弧が削除されています)。

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

実際のシナリオに現れる可能性があるこのような型のフェデレーテッド値の 1 例として、サーバーがフェデレーテッドトレーニングに参加するデバイスのグループにブロードキャストする、ハイパーパラメータ(学習率、クリッピングノルムなど)があります。

別の例としては、サーバーで事前トレーニングされた機械学習モデルのパラメータのセットがあります。これはクライアントデバイスのグループにブロードキャストされ、そこで各ユーザーに合わせてパーソナライズすることができます。

例えば、単純な 1 次元線形回帰モデルのfloat32パラメータのペアabがあるとします。TFF で使用するには、このようなモデルの(非フェデレーテッドの)型は以下のように構築します。出力された型文字列の山括弧<>は、名前付きタプルまたは名前なしタプルのコンパクトな TFF 表記です。

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

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

上記ではdtypeのみを指定していることに注意してください。非スカラー型もサポートされています。上のコードでは、tf.float32はより一般的なtff.TensorType(dtype=tf.float32, shape=[])のショートカット表記です。

このモデルがクライアントにブロードキャストされると、結果のフェデレーテッド値の型は以下のように表すことができます。

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

上記のフェデレーテッド float との対称性から、このような型をフェデレーテッドタプルと呼びます。より一般的には、メンバ構成要素が XYZ であるフェデレーテッド値を指す場合にフェデレーテッド XYZ という名称をよく使用します。そのため、フェデレーテッドタプルフェデレーテッドシーケンスフェデレーテッドモデルなどについて説明しています。

さて、float32@CLIENTSの話に戻ります。これは複数のデバイス間で複製されているように見えますが、すべてのメンバが同じなので、実際には単一のfloat32です。一般的には、すべて等しいフェデレーテッド型、つまり T@G 形式のものは、非フェデレーテッド型の T と同型であると考えることができます。それはどちらの場合でも、実際にはT型のアイテムが 1 つだけ(複製される可能性はありますが)存在するからです。

TT@G間の同型性を考えると、後者の型がどのような目的で役に立つのか疑問に思われるかもしれません。先をお読みください。

配置

設計の概要

前のセクションでは、配置、すなわちフェデレーテッド値を共同でホストする可能性のあるシステムの参加者のグループの概念について、また配置の仕様の例としてtff.CLIENTSの使用について紹介しました。

なぜ配置という概念が非常に基本的なものであり、TFF 型システムに組み込む必要があるのかを説明するにあたり、本チュートリアルの冒頭で述べた TFF の使用目的について思い出してみてください。

本チュートリアルでは、TFF コードはシミュレートされた環境でローカルで実行するだけですが、TFF が目標としているのは、分散システム内の物理デバイスのグループ(Android を実行しているモバイルデバイスや組み込みデバイスを含む可能性があります)にデプロイして実行できるようなコード記述ができるようにすることです。これらのデバイスは、それぞれのシステムで果たす役割(エンドユーザーデバイス、集中型コーディネータ、多階層アーキテクチャの中間レイヤーなど)に応じて、ローカルで実行する個別の命令セットを受け取ります。デバイスのどのサブセットがどのようなコードを実行し、異なる部分のデータを物理的にどこでマテリアライズするかについて推論ができることは重要です。

これは、例えばモバイルデバイス上のアプリケーションデータなどを扱う場合、特に重要です。データは非公開で機密性が高い可能性があるため、このデータがデバイス外に出ることがないよう静的に検証する(およびデータの処理方法に関する事実を証明する)機能が必要です。配置仕様は、これをサポートするために設計された仕組みの 1 つです。

TFF はデータを中心としたプログラミング環境として設計されているため、演算子およびその演算子がどこで実行されるかに焦点を当てた既存のフレームワークとは異なり、TFF はデータ、そのデータがマテリアライズされる場所、およびその変換方法に焦点を当てています。その結果、TFF では配置はデータの演算子のプロパティとしてではなく、データのプロパティとしてモデル化されます。実際、次のセクションで説明するように、一部の TFF 演算子は複数の場所にまたがり、単一のマシンやマシンのグループで実行されるのではなく、いわば「ネットワーク内」で実行されます。

特定の値の型を(単にTとして表すのとは対照的に)T@Gまたは{T}@Gとして表すことは、データ配置の決定を明示的にするとともに、TFF で記述されたプログラムの静的解析と合わせて、デバイス上の機密性の高いデータに正式なプライバシー保証を提供する基盤となります。

ただし、この時点で注意すべき重要な点は、TFF ユーザーにはデータ(配置)をホストする参加デバイスの グループ を明示的に示すよう推奨していますが、プログラマが個々の参加者の生データや身元情報を取り扱うことは決してないということです。

(注意: 本チュートリアルの範囲外になりますが、上記には注目すべき例外が 1 つあることに触れておきます。tff.federated_collect演算子は、低レベルのプリミティブとして、特殊な状況に限った使用を意図しています。これを回避可能な状況で明示的に使用することは、この先使用する可能性のあるアプリケーションを制限してしまう可能性があるため、推奨できません。例えば、静的解析の過程で任意の計算がそのような低レベルの仕組みを使用していると判断した場合、特定のタイプのデータへのアクセスを許可しない場合があります。)

設計上、TFF コードの本体内ではtff.CLIENTSで表されるグループを構成するデバイスを列挙したり、グループ内に特定のデバイスの存在するかどうかを調べたりする方法がありません。フェデレーテッドコア API、基礎となるアーキテクチャの抽象化セット、シミュレーションをサポートするために提供するコア ランタイム インフラストラクチャには、デバイスやクライアントを識別する概念がありません。記述するすべての計算ロジックは、クライアントグループ全体に対する演算子として表現されます。

ここで、フェデレーテッド型の値が Python のdictとは異なり、メンバ構成要素を単純に列挙することはできないと先に述べました。TFF プログラムのロジックが扱う値は、個々の参加者ではなく、配置(グループ)に関連付けられると考えてください。

TFF においても、配置はすべて第一級オブジェクトとなるように設計されており、placement型(API ではtff.PlacementTypeで表現)のパラメータおよび結果として表示できます。今後は配置の変換や結合を可能にする多様な演算子の提供を予定していますが、それは本チュートリアルの範囲外です。現時点では、intboolが Python の不透明な組み込み型であるのと同様に、placementは TFF の不透明なプリミティブ組み込み型であると考えれば十分です。これはtff.CLIENTSがこの型の定数リテラルであり、1int型の定数リテラルであることに似ています。

配置を指定する

TFF はtff.CLIENTStff.SERVERという 2 つの基本的な配置リテラルを提供し、クライアント/サーバーアーキテクチャとして、自然にモデル化された多様な実用的シナリオを簡単に表現できるようにしています。複数のクライアントデバイス(携帯電話、組み込みデバイス、分散データベース、センサーなど)を使用して、単一の集中型サーバーコーディネータでオーケストレーションします。TFF はカスタム配置、複数のクライアントグループ、多層化、その他より一般的な分散アーキテクチャもサポートできるように設計されていますが、それらに関する説明は本チュートリアルの範囲外となります。

TFF はtff.CLIENTStff.SERVERが実際に何を表すかを規定していません。

特に、tff.SERVERは単一の物理デバイス(シングルトングループのメンバ)である場合がありますが、ステートマシンレプリケーションを実行しているフォールト トレラント クラスタ内のレプリカのグループである場合もあります。むしろ、前のセクションで述べたall_equalの部分を使用して、通常サーバーでは単一のデータアイテムのみを処理するという事実を表現します。

同様に、一部のアプリケーションでは、tff.CLIENTSはシステム内のすべてのクライアントを表す場合があり、フェデレーテッドラーニングの文脈では集団と呼ぶことがあります。しかし、例えばフェデレーテッドアベレージングのプロダクション実装では、トレーニングの特定のラウンドに参加するために選択されたクライアントのサブセットである cohort を表します。抽象的に定義された配置は、それらが出現する計算が実行のためにデプロイされるとき(または本チュートリアルで実演しているように、シミュレートされた環境で Python 関数のように単純に呼び出されるとき)に具体的な意味が与えられます。このローカルシミュレーションでは、入力として供給されたフェデレーテッドデータによってクライアントのグループを決定します。

フェデレーテッド計算

フェデレーテッド計算を宣言する

TFF は、モジュール開発をサポートする、強力に型付けされた関数型プログラミング環境として設計されています。

TFF の構成の基本的な単位はフェデレーテッド計算、すなわちフェデレーテッド値を入力として受け入れ、フェデレーテッド値を出力として返すロジックのセクションです。ここでは、前述の例のセンサーアレイから報告された温度の平均値を算出する計算の定義方法を以下に示します。

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

上記のコードを見るとこの時点では、TensorFlow にはtf.functionのような合成可能な単位を定義するデコレータ構造はすでに存在していないのではないか、もし存在するならば、なぜ別のデコレータの構造を導入するのか、そしてどう違うのか、といった疑問が生じるかもしれません。

その答えを簡単に言うと、tff.federated_computationラッパーによって生成されたコードは TensorFlow でも Python でもないということです。これは、内部プラットフォームに依存しないグルー言語による分散システムの仕様です。この時点では間違いなく不可解に聞こえますが、このように直感的に解釈されたフェデレーテッド計算は分散システムの抽象的な仕様であることを頭に入れておいてください。この後すぐに説明をします。

まず最初に、定義を少し入力してみましょう。通常 TFF 計算は、パラメータの有無に関わらず明確に定義された型シグネチャを使用して、関数としてモデル化されます。以下に示すようにtype_signatureプロパティをクエリすると、計算の型シグネチャを出力することができます。

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

この型シグネチャは、計算がクライアントデバイス上の様々なセンサーの測定値のコレクションを受け入れ、サーバー上で単一の平均値を返すことを示しています。

この計算の入力と出力は異なる場所CLIENTS上とSERVER上)にあります。先に進む前に、これについて少し考えてみましょう。前のセクションで配置に関して述べた、TFF の操作が場所をまたいでネットワーク内でどのように実行されるかについて、そして先程説明した、フェデレーテッド計算が分散システムの抽象的な仕様を表すものだということについて思い出してください。ここではそのような計算の 1 つを、つまりデータをクライアントデバイスで消費して集約結果をサーバで得る単純な分散システムを定義したにすぎません。

多くの実用的なシナリオにおいて、トップレベルのタスクを表す計算はサーバーに入力を受け入れ、サーバーに出力を報告する傾向があります。これはサーバーを起点と終点にしたクエリによって、計算がトリガされる可能性があるという考えを反映しています。

ただし、FC API はこの仮定を課さないため、内部で使用するビルディングブロックの多く(API 内にある数多くのtff.federated_...演算子を含む)には配置の異なる入力と出力があります。そのため通常は、フェデレーテッド計算をサーバー上で実行するもの、あるいはサーバーが実行するものと考えてはいけません。サーバーは、フェデレーテッド計算の参加者の 1 タイプにすぎません。このような計算の仕組みを考える際には、単一の集中型コーディネータの視点ではなく、常にグローバルなネットワーク全体の視点をデフォルトで考えるのがベストです。

一般に、関数型シグネチャは、入力と出力の型TUに対して、それぞれ(T -> U)のようにコンパクトに表現します。デコレータの引数には、フォーマルパラメータの型(この場合はsensor_readingsなど)を指定します。結果の型を指定する必要はなく、これは自動的に決定されます。

TFF が提供するポリモーフィズムの形式は限られていますが、コードのプロパティの理解、デバッグ、および正式な検証を容易にするために、プログラマの皆さんには扱うデータの型を明示的に指定することを強くお勧めします。場合によっては型の明示的な指定が必要条件であることもあります(例えば現時点では、ポリモーフィック計算の直接実行はできません)。

フェデレーテッド計算を実行する

TFF は開発やデバッグをサポートするために、以下に示すように、この方法で定義された計算を Python の関数として直接呼び出すことができます。計算がall_equalの部分をFalse設定にしたフェデレーテッド型の値を期待している場合、Python では単純なlistとして与えることができます。また、all_equalの部分をTrue設定にしたフェデレーテッド型の場合は、(単一の)メンバ構成要素を直接与えることができます。これは結果を報告する方法でもあります。

get_average_temperature([68.5, 70.3, 69.8])
69.53334

このような計算をシミュレーションモードで実行する場合、あなたにはネットワーク内の任意の場所で入力を供給して出力を使用することができる、システム全体のビューを持った外部オブザーバーとしての役割があります。ここでは入力時にクライアントの値を供給し、サーバーの結果を使用しました。

ここで、先ほどのグルー言語でコードを出力するtff.federated_computationデコレータに関して保留していた説明に戻りましょう。TFF 計算のロジックは Python でも(上記のようにtff.federated_computationで装飾するだけで)普通の関数として表現できます。このノートブックにある他の Python 関数と同様に Python の引数を使用して直接呼び出すこともできますが、その裏では、先にも述べたように TFF 計算は実は Python ではありません

これが何を意味するかというと、Python インタプリタはtff.federated_computationで装飾された関数に遭遇すると、その関数の本体内のステートメントを一度だけ(定義時に)トレースします。そして実行目的や別の計算にサブコンポーネントとして組み込む目的で後に使用するために、計算ロジックのシリアライズされた表現を構築するということです。

以下のように print 文を追加すると、これを確認することができます。

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

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

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

フェデレーテッド計算を定義する Python のコードは、非 Eager なコンテキストで TensorFlow グラフを構築する Python のコードの考え方に似ています (TensorFlow の非 Eager な使用法をご存知でない場合には、後で実行される演算のグラフを定義した Python のコードは、実際にはその場で実行されるわけではないと考えてみてください)。TensorFlow で非 Eager なグラフ構築のコードは Python ですが、このコードによって構築された TensorFlow のグラフはプラットフォームに依存しないので、シリアライズが可能です。

同様に、TFF の計算は Python で定義しますが、先ほど示した例のtff.federated_meanなど、その本体内の Python ステートメントは、内部でポータブルかつプラットフォームに依存しないシリアライズ可能な表現にコンパイルされます。

開発者が直接この表現を使って作業をすることはないので、この表現の詳細を気にする必要はありませんが、TFF 計算は基本的に非 Eager であり、任意の Python の状態をキャプチャできないという事実に注意してください。tff.federated_computationで装飾された Python 関数の本体をシリアライズ前するにトレースする際、TFF 計算の本体に含まれる Python のコードを定義時に実行します。呼び出し時に再びトレースすることはありません。(ポリモーフィックな場合を除きます。詳細についてはドキュメントのページをご覧ください。)

なぜ Python ではない専用の内部表現を導入したのか、疑問に思われるかもしれません。その理由の 1 つは、最終的には TFF の計算を実際の物理的環境に展開し、Python を使用できないモバイルデバイスや組み込みデバイスでもホストされることを想定しているからです。

もう 1 つの理由は、個々の参加者のローカルな振る舞いを表現する Python プログラムとは対照的に、TFF 計算は分散システムのグローバルな振る舞いを表現するからです。上記の単純な例では、特殊な演算子tff.federated_meanを使用してクライアントデバイス上のデータを受け取り、結果をサーバーに送信することが分かります。

演算子tff.federated_meanはローカルで実行せず、先に述べたように複数のシステム参加者の動作を調整する分散システムを表現しているため、Python の一般的な演算子として簡単にモデル化することはできません。このような演算子をフェデレーテッド演算子と呼び、Python の一般的(ローカル)な演算子とは区別します。

TFF 型システムや TFF の言語でサポートされる演算子の基本的なセットは Python の演算子と大きく異なるため、専用の表現を使用する必要があるのです。

フェデレーテッド計算を作成する

上で述べたように、フェデレーテッド計算とその構成要素は、分散システムのモデルとして最もよく理解されており、フェデレーテッド計算の作成は、単純な分散システムからより複雑な分散システムを作成することであると考えることができます。演算子tff.federated_meanは、型シグネチャ({T}@CLIENTS -> T@SERVER)を持つ、ある種の組み込みテンプレートのフェデレーテッド計算と考えることができます(実際、計算の作成と同様に、この演算子の構造も複雑で、それをより単純な演算子に内部で分解します)。

フェデレーテッド計算を作成する場合も同様です。計算get_average_temperatureは、tff.federated_computationで装飾された別の Python 関数の本体内で呼び出すことができます。そうすると、前にtff.federated_meanがその本体内に埋め込まれたのと同じような方法で、親の本体内に埋め込まれます。

注意すべき重要な制限事項は、tff.federated_computationで装飾された Python 関数の本体は、フェデレーテッド演算子のみで構成する必要があるいうことです。つまり、直接 TensorFlow 演算子を含むことはできません。例えば、直接tf.nestインターフェースを使用してフェデレーテッド値のペアを追加することはできません。TensorFlow コードの使用は、次のセクションで説明するtff.tf_computationで装飾されたコードのブロックに制限する必要があります。そのようにしてラップされた場合にのみ、tff.federated_computationの本体内でラップされた TensorFlow コードを呼び出すことができます。

この分離の理由には、技術的な理由(非テンソルでtf.addのような演算子が使えるように仕向けるのは困難であること)およびアーキテクチャ的な理由があります。フェデレーテッド計算の言語(つまりtff.federated_computationで装飾された Python 関数のシリアライズされた本体から構築したロジック)は、プラットフォームに依存しないグルー言語として機能するように設計されています。このグルー言語は現在、TensorFlow コードの埋め込みセクション(tff.tf_computationブロックに限定)から分散システムを構築するために使用されています。将来的には、入力パイプラインを表すリレーショナル データベース クエリのような TensorFlow 以外のロジックのセクションを埋め込んで、同じグルー言語(tff.federated_computationブロック)を使用してこれらすべてを接続できるようにする必要があると予想しています。

TensorFlow のロジック

TensorFlow 計算を宣言する

TFF は TensorFlow で使用するように設計されています。したがって、TFF で記述するコードの大部分は通常の(つまりローカルで実行する)TensorFlow のコードになるはずです。上で述べたように、tff.tf_computationで装飾するだけで、TensorFlow のコードを TFF で使用できるようになります。

例えば、数値を受け取り、それに0.5を加える関数を実装する方法を以下に示します。

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

もう一度これを見ると、単純にtf.functionなどの既存の仕組みを使用するのではなく、なぜ別のデコレータtff.tf_computationを定義する必要があるのか疑問に思うかもしれません。前のセクションとは異なり、ここでは TensorFlow コードの普通のブロックを扱っています。

これにはいくつかの理由があります。その全体的な扱いに関しては本チュートリアルの範囲外になってしまいますが、ここに主要な理由を挙げておきます。

  • TensorFlow コードを使用して実装した再利用可能なビルディングブロックをフェデレーテッド計算の本体に埋め込むには、定義時にトレースされてシリアライズされる、型シグネチャを持つなどの特定のプロパティを満たす必要があります。これには通常、何らかの形のデコレータが必要です。

通常は、可能な限りtf.functionのような TensorFlow のネイティブメカニズムの使用を推奨しています。これは TFF のデコレータが Eager 関数と相互作用するこの確実な方法は、進化が期待できるからです。

ここで、上記のコードスニペットの例に戻りますが、先ほど定義した計算add_halfは、他の TFF 計算と同様に TFF で処理することができます。特に、これには TFF 型シグネチャがあります。

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

この型シグネチャには配置がないことに注意してください。TensorFlow の計算は、フェデレーテッド型を消費または返すことができません。

また、add_halfを他の計算のビルディングブロックとして使用することもできます。例として、tff.federated_map 演算子を使用して、クライアントデバイス上のフェデレーテッド float のすべてのメンバ構成要素に点ごとにadd_halfを適用する方法を以下に示します。

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

TensorFlow 計算を実行する

tff.tf_computationで定義された計算の実行は、tff.federated_computationで説明したのと同じルールに従います。これらは通常の Python の callable として以下のように呼び出すことができます。

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

繰り返しますが、この方法でadd_half_on_clientsの計算を呼び出すと、分散プロセスをシミュレートすることに注目してください。データをクライアント上で消費し、クライアント上で返します。実際、この計算は各クライアントにローカルアクションを実行させます。(実際にはこのようなプロセスをオーケストレーションする際に必要になる可能性はありますが、)このシステム内ではtff.SERVERを明示的に言及していません。この方法で定義された計算は、MapReduceMapステージの概念に類似していると考えてください。

また、前のセクションで、TFF の計算を定義時にシリアライズすることについて言及しましたが、これはtff.tf_computationのコードにも該当します。add_half_on_clientsの Python 本体は定義時に一度だけトレースします。それ以降の呼び出しでは、TFF はそのシリアライズされた表現を使用します。

tff.federated_computationで装飾された Python メソッドとtff.tf_computationで装飾された Python メソッドの唯一の違いは、(前者が TensorFlow コードを直接埋め込むことが許されないのに対し)後者は TensorFlow のグラフとしてシリアライズされることです。

内部的にtff.tf_computationで装飾された各メソッドは、計算の構造をキャプチャできるようにするため、一時的に Eager execution を無効化します。Eager execution をローカルで無効化し、正しくシリアライズできるように計算のロジックを記述している限りは、Eager な TensorFlow、AutoGraph、TensorFlow 2.0 コンストラクトなどを使用しても問題ありません。

例えば、以下のコードは失敗します。

try:

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

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

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

上記では、tff.tf_computationadd_tenの本体で内部的に構築するグラフの外側で、シリアル化のプロセス中に既にconstant_10が構築されるため、失敗します。

その一方で、内部にtff.tf_computationを呼び出す際に、現在のグラフを変更する python 関数を呼び出すことについては問題ありません。

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

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

add_ten(5.0)
15.0

TensorFlow のシリアル化の仕組みが進化をしているため、TFF の計算のシリアライズ方法の詳細についても今後進化していくと予想しています。

tf.data.Datasetで作業する

前述したように、tff.tf_computationのユニークな特徴は、正式なパラメータとして抽象的に定義したtf.data.Datasetをコードを扱うことができるということです。TensorFlow でデータセットとして表現するパラメータは、tff.SequenceTypeコンストラクタを使用して宣言する必要があります。

例えば、型の仕様tff.SequenceType(tf.float32)は、TFF の float 要素の抽象シーケンスを定義します。シーケンスにはテンソルまたは複雑な入れ子構造のいずれかを含めることができます。(これらの例については後ほど説明します。)T型のアイテムのシーケンスを簡潔に表現すると、T*となります。

float32_sequence = tff.SequenceType(tf.float32)

str(float32_sequence)
'float32*'

温度センサーの例では、各センサーが単一の温度測定値ではなく、それぞれ複数の温度測定値を保持しているとします。ここではtf.data.Dataset.reduce演算子を使用して、TensorFlow で 単一のローカルデータセットの平均温度を計算する TFF 計算の定義方法を説明します。

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

tff.tf_computationで装飾されたメソッドの本体では、TFF のシーケンス型の形式的なパラメータは、tf.data.Datasetのように振る舞うオブジェクトとして単純に表現されます。つまり、同じプロパティとメソッドをサポートします。(現在はその型のサブクラスとして実装されていませんが、これは TensorFlow のデータセットのサポートの進化に伴って変更される可能性があります。)

以下のようにして簡単に確認することができます。

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

foo([1, 2, 3])
6

一般的なtf.data.Datasetとは異なり、これらのデータセットのようなオブジェクトはプレースホルダーであることに注意してください。これらのオブジェクトは抽象的なシーケンス型のパラメータを表現しているため、要素は一切含んでおらず、具体的なコンテキストで使用する際に具体的なデータにバインドされます。抽象的に定義されたプレースホルダー データセットのサポートについては、現時点ではまだ多少限界があり TFF の初期段階なので、ある種の制限に遭遇する可能性があります。しかし、本チュートリアルで心配する必要はありません(詳細についてはドキュメントのページをご覧ください)。

本チュートリアルのように、シーケンスを受け入れる計算をシミュレーションモードでローカルで実行する場合、以下のように Python のリストとしてシーケンスを与えることができます(他の方法、例えば Eager モードのtf.data.Datasetでも可能ですが、今のところはシンプルにしておきます)。

get_local_temperature_average([68.5, 70.3, 69.8])
69.53333

他のすべての TFF 型もそうですが、上で定義したようなシーケンスは、tff.StructTypeコンストラクタを使用して入れ子構造を定義することができます。例えば、ABのペアのシーケンスを受け入れ、それらの積の和を返す計算を宣言する方法を以下に示します。TFF の型シグネチャがどのようにデータセットのoutput_typesoutput_shapesに変換しているかを確認できるよう、計算の本文にトレース文を入れています。

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

tf.data.Datasetsを形式的なパラメータとして使用するためのサポートについては、本チュートリアルで使用されているような単純なシナリオでは機能しますが、まだ多少制限があり、進化を続けているところです。

すべてを統合する

さて、ここで再びフェデレーテッド設定で TensorFlow 計算を使用してみましょう。各センサーが温度測定値のローカルシーケンスを持つセンサー群があるとします。以下のようにセンサーのローカル平均値を平均化して、グローバルな平均温度を計算することができます。

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

これは全てのクライアントの全てのローカルの温度測定値にわたる単純平均値ではないことに注意してください。単純平均値には、各クライアントがローカルで保持する測定値の数に従い、それぞれの測定値に対する重み付けが必要となるからです。これは上記コードの更新の練習として、読者の皆さんに残しておくことにします。tff.federated_mean演算子は、重みをオプションの第 2 引数(フェデレーテッド float が予想される)として受け入れます。

さらに、get_global_temperature_averageの入力がフェデレーテッド float シーケンスになっていることにも注意してください。フェデレーテッドシーケンスはフェデレーテッドラーニングでデバイス上のデータを一般的に表現する方法であり、シーケンス要素は一般的にデータのバッチを表現します(後の例でご覧になれます)。

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

ここで Python のデータのサンプルに対してローカルで計算を実行する方法を説明します。入力を供給する方法は、listlistとなっていることに注意してください。外側のリストはtff.CLIENTSで表現されるグループ内のデバイスをイテレートし、内側のリストは各デバイスのローカルシーケンス内の要素をイテレートします。

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

これでチュートリアルの第 1 部は終了です。次は第 2 部にお進みください。