Copyright 2020 The TensorFlow Authors.
TensorFlow の NumPy API
概要
TensorFlow では、tf.experimental.numpy
を利用してNumPy API のサブセットを実装します。これにより、TensorFlow により高速化された NumPy コードを実行し、TensorFlow のすべて API にもアクセスできます。
セットアップ
NumPy 動作の有効化
tnp
を NumPy として使用するには、TensorFlow の NumPy の動作を有効にしてください。
この呼び出しによって、TensorFlow での型昇格が可能になり、リテラルからテンソルに変換される場合に、型推論も Numpy の標準により厳格に従うように変更されます。
注意: この呼び出しは、tf.experimental.numpy
モジュールだけでなく、TensorFlow 全体の動作を変更します。
TensorFlow NumPy ND 配列
ND 配列と呼ばれる tf.experimental.numpy.ndarray
は、特定のデバイスに配置されたある dtype
の多次元の密な配列を表します。tf.Tensor
のエイリアスです。ndarray.T
、ndarray.reshape
、ndarray.ravel
などの便利なメソッドについては、ND 配列クラスをご覧ください。
まず、ND 配列オブジェクトを作成してから、さまざまなメソッドを呼び出します。
型昇格
TensorFlow の型昇格には 4 つのオプションがあります。
デフォルトでは、混合型の演算に対し、TensorFlow は型を昇格する代わりにエラーを発します。
tf.numpy.experimental_enable_numpy_behavior()
を実行すると、NumPy
型の昇格ルール(以下に説明)を使うように TensorFlow が切り替えられます。TensorFlow 2.15 以降で使用できる以下の新しいオプションが 2 つあります(詳細は、TF NumPy 型昇格をご覧ください)。
tf.numpy.experimental_enable_numpy_behavior(dtype_conversion_mode="all")
は、Jax の型昇格ルールを使用します。tf.numpy.experimental_enable_numpy_behavior(dtype_conversion_mode="safe")
は Jax の型昇格ルールを使用しますが、特定の安全でない昇格を許可しません。
NumPy 型昇格
TensorFlow NumPy API には、リテラルを ND 配列に変換するためと ND 配列入力で型昇格を実行するための明確に定義されたセマンティクスがあります。詳細については、np.result_type
をご覧ください。
TensorFlow API は tf.Tensor
入力を変更せずそのままにし、それに対して型昇格を実行しませんが、TensorFlow NumPy API は NumPy 型昇格のルールに従って、すべての入力を昇格します。次の例では、型昇格を行います。まず、さまざまな型の ND 配列入力で加算を実行し、出力の型を確認します。これらの型昇格は、TensorFlow API では行えません。
最後に、ndarray.asarray
を使ってリテラルをND 配列に変換し、結果の型を確認します。
リテラルを ND 配列に変換する際、NumPy は tnp.int64
や tnp.float64
といった幅広い型を優先します。一方、tf.convert_to_tensor
は、tf.int32
と tf.float32
の型を優先して定数を tf.Tensor
に変換します。TensorFlow NumPy API は、整数に関しては NumPy の動作に従っています。浮動小数点数については、experimental_enable_numpy_behavior
の prefer_float32
引数によって、tf.float64
よりも tf.float32
を優先するかどうかを制御することができます(デフォルトは False
です)。以下に例を示します。
ブロードキャスト
TensorFlow と同様に、NumPy は「ブロードキャスト」値の豊富なセマンティクスを定義します。詳細については、NumPy ブロードキャストガイドを確認し、これを TensorFlow ブロードキャストセマンティクスと比較してください。
インデックス作成
NumPy は、非常に洗練されたインデックス作成ルールを定義しています。NumPy インデックスガイドを参照してください。以下では、インデックスとして ND 配列が使用されていることに注意してください。
サンプルモデル
次に、モデルを作成して推論を実行する方法を見てみます。この簡単なモデルは、relu レイヤーとそれに続く線形射影を適用します。後のセクションでは、TensorFlow のGradientTape
を使用してこのモデルの勾配を計算する方法を示します。
TensorFlow NumPy および NumPy
TensorFlow NumPy は、完全な NumPy 仕様のサブセットを実装します。シンボルは、今後追加される予定ですが、近い将来にサポートされなくなる体系的な機能があります。これらには、NumPy C API サポート、Swig 統合、Fortran ストレージ優先順位、ビュー、stride_tricks
、およびいくつかのdtype
(np.recarray
や np.object
)が含まれます。詳細については、 TensorFlow NumPy API ドキュメントをご覧ください。
NumPy 相互運用性
TensorFlow ND 配列は、NumPy 関数と相互運用できます。これらのオブジェクトは、__array__
インターフェースを実装します。NumPy はこのインターフェースを使用して、関数の引数を処理する前にnp.ndarray
値に変換します。
同様に、TensorFlow NumPy 関数は、np.ndarray
などのさまざまなタイプの入力を受け入れることができます。これらの入力は、ndarray.asarray
を呼び出すことにより、ND 配列に変換されます。
ND 配列をnp.ndarray
との間で変換すると、実際のデータコピーがトリガーされる場合があります。詳細については、バッファコピーのセクションを参照してください。
バッファコピー
TensorFlow NumPy を NumPy コードと混在させると、データコピーがトリガーされる場合があります。これは、TensorFlow NumPy のメモリアライメントに関する要件が NumPy の要件よりも厳しいためです。
np.ndarray
が TensorFlow Numpy に渡されると、アライメント要件を確認し、必要に応じてコピーがトリガーされます。ND 配列 CPU バッファを NumPy に渡す場合、通常、バッファはアライメント要件を満たし、NumPy はコピーを作成する必要はありません。
ND 配列は、ローカル CPU メモリ以外のデバイスに配置されたバッファを参照できます。このような場合、NumPy 関数を呼び出すと、必要に応じてネットワークまたはデバイス全体でコピーが作成されます。
このため、NumPy API 呼び出しとの混合は通常、注意して行い、ユーザーはデータのコピーのオーバーヘッドに注意する必要があります。TensorFlow NumPy 呼び出しを TensorFlow 呼び出しとインターリーブすることは一般的に安全であり、データのコピーを避けられます。 詳細については、TensorFlow の相互運用性のセクションをご覧ください。
演算子の優先順位
TensorFlow NumPy は、NumPy よりも優先順位の高い__array_priority__
を定義します。つまり、ND 配列とnp.ndarray
の両方を含む演算子の場合、前者が優先されます。np.ndarray
入力は ND 配列に変換され、演算子の TensorFlow NumPy 実装が呼び出されます。
TF NumPy と TensorFlow
TensorFlow NumPy は TensorFlow の上に構築されているため、TensorFlow とシームレスに相互運用できます。
tf.Tensor
と ND 配列
ND 配列は tf.Tensor
のエイリアスであるため、実際のデータのコピーを呼び出さずに混在させることが可能です。
TensorFlow 相互運用性
ND 配列は tf.Tensor
のエイリアスにすぎないため、TensorFlow API に渡すことができます。前述のように、このような相互運用では、アクセラレータやリモートデバイスに配置されたデータであっても、データのコピーは行われません。
逆に言えば、tf.Tensor
オブジェクトを、データのコピーを実行せずに、tf.experimental.numpy
API に渡すことができます。
勾配とヤコビアン: tf.GradientTape
TensorFlow の GradientTape は、TensorFlow と TensorFlow NumPy コードを介してバックプロパゲーションに使用できます。
サンプルモデルセクションで作成されたモデルを使用して、勾配とヤコビアンを計算します。
トレースコンパイル: tf.function
Tensorflow の tf.function
は、コードを「トレースコンパイル」し、これらのトレースを最適化してパフォーマンスを大幅に向上させます。グラフと関数の概要を参照してください。
また、tf.function
を使用して、TensorFlow NumPy コードを最適化することもできます。以下は、スピードアップを示す簡単な例です。tf.function
コードの本文には、TensorFlow NumPy API への呼び出しが含まれていることに注意してください。
ベクトル化:tf.vectorized_map
TensorFlow には、並列ループのベクトル化のサポートが組み込まれているため、10 倍から 100 倍のスピードアップが可能です。これらのスピードアップは、tf.vectorized_map
API を介して実行でき、TensorFlow NumPy にも適用されます。
w.r.t. (対応する入力バッチ要素)バッチで各出力の勾配を計算すると便利な場合があります。このような計算は、以下に示すように tf.vectorized_map
を使用して効率的に実行できます。
デバイスに配置する
TensorFlow NumPy は、CPU、GPU、TPU、およびリモートデバイスに演算を配置できます。デバイスにおける配置には標準の TensorFlow メカニズムを使用します。以下の簡単な例は、すべてのデバイスを一覧表示してから、特定のデバイスに計算を配置する方法を示しています。
ここでは取り上げませんが、TensorFlow には、デバイス間で計算を複製し、集合的な削減を実行するための API もあります。
デバイスをリストする
使用するデバイスを見つけるには、tf.config.list_logical_devices
およびtf.config.list_physical_devices
を使用します。
演算の配置:tf.device
デバイスに演算を配置するには、tf.device
スコープでデバイスを呼び出します。
デバイス間での ND 配列のコピー: tnp.copy
特定のデバイススコープで tnp.copy
を呼び出すと、データがそのデバイスに既に存在しない限り、そのデバイスにデータがコピーされます。
パフォーマンスの比較
TensorFlow NumPy は、CPU、GPU、TPU にディスパッチできる高度に最適化された TensorFlow カーネルを使用します。TensorFlow は、演算の融合など、多くのコンパイラ最適化も実行し、パフォーマンスとメモリを向上します。詳細については、Grappler を使用した TensorFlow グラフの最適化をご覧ください。
ただし、TensorFlow では、NumPy と比較してディスパッチ演算のオーバーヘッドが高くなります。小規模な演算(約 10 マイクロ秒未満)で構成されるワークロードの場合、これらのオーバーヘッドがランタイムを支配する可能性があり、NumPy はより優れたパフォーマンスを提供する可能性があります。その他の場合、一般的に TensorFlow を使用するとパフォーマンスが向上するはずです。
以下のベンチマークを実行して、さまざまな入力サイズでの NumPy と TensorFlow Numpy のパフォーマンスを比較します。
参考資料
TensorFlow NumPy: 分散型画像分類のチュートリアル
TensorFlow NumPy: Keras と分散ストラテジー