概要
このガイドでは、TensorFlow.js でカスタム オペレーション (ops)、カーネル、およびグラデーションを定義するメカニズムの概要を説明します。主要な概念の概要と、実際の概念を示すコードへのポインタを提供することを目的としています。
このガイドは誰を対象としていますか?
これは、TensorFlow.js の内部の一部に触れたかなり高度なガイドであり、特に次のグループに役立つ可能性があります。
- さまざまな数学的演算の動作のカスタマイズに興味のある TensorFlow.js の上級ユーザー (例: 既存の勾配実装をオーバーライドする研究者、またはライブラリで欠落している機能にパッチを適用する必要があるユーザー)
- TensorFlow.js を拡張するライブラリを構築するユーザー (たとえば、TensorFlow.js プリミティブまたは新しい TensorFlow.js バックエンド上に構築された一般的な線形代数ライブラリ)。
- tensorflow.js に新しい操作を提供することに興味があり、これらのメカニズムがどのように機能するかの概要を知りたいユーザー。
これは、TensorFlow.js が内部実装メカニズムに関わるものであるため、TensorFlow.js の一般的な使用方法に関するガイドではありません。 TensorFlow.js を使用するためにこれらのメカニズムを理解する必要はありません
このガイドを最大限に活用するには、TensorFlow.js ソース コードを読むことに慣れている (または喜んで読んでみる) 必要があります。
用語
このガイドでは、いくつかの重要な用語を前もって説明するのに役立ちます。
演算 (Ops) — 1 つまたは複数のテンソルを出力として生成する、1 つまたは複数のテンソルに対する数学的演算。演算は「高レベル」コードであり、他の演算を使用してロジックを定義できます。
カーネル— 特定のハードウェア/プラットフォームの機能に関連付けられたオプの特定の実装。カーネルは「低レベル」でバックエンド固有です。一部の操作は操作からカーネルへ 1 対 1 のマッピングを持ちますが、他の操作は複数のカーネルを使用します。
Gradient / GradFunc — 入力に関する関数の導関数を計算する演算/カーネルの「バックワード モード」定義。グラデーションは「高レベル」コード (バックエンド固有ではない) であり、他の操作やカーネルを呼び出すことができます。
カーネル レジストリ- (カーネル名、バックエンド名)タプルからカーネル実装へのマップ。
グラデーション レジストリ—カーネル名からグラデーション実装へのマップ。
コード構成
カーネルはバックエンド固有であり、それぞれのバックエンド フォルダー (例: tfjs-backend-cpu ) で定義されます。
カスタム演算、カーネル、およびグラデーションをこれらのパッケージ内で定義する必要はありません。ただし、実装では同様のシンボルが使用されることがよくあります。
カスタム Ops の実装
カスタム op を考える 1 つの方法は、多くの場合、テンソルを入力として使用して、何らかのテンソル出力を返す JavaScript 関数として考えることです。
- 一部の操作は既存の操作に関して完全に定義でき、これらの関数を直接インポートして呼び出すだけで済みます。以下に例を示します。
- op の実装は、バックエンド固有のカーネルにディスパッチすることもできます。これは
Engine.runKernelを介して行われ、「カスタム カーネルの実装」セクションで詳しく説明されます。以下に例を示します。
カスタム カーネルの実装
バックエンド固有のカーネル実装により、特定の操作のロジックの最適化された実装が可能になります。カーネルはtf.engine().runKernel()を呼び出す操作によって呼び出されます。カーネル実装は 4 つの要素によって定義されます。
- カーネル名。
- カーネルが実装されているバックエンド。
- 入力: カーネル関数への Tensor 引数。
- 属性: カーネル関数への非テンソル引数。
以下はカーネル実装の例です。実装に使用される規則はバックエンド固有であり、各特定のバックエンドの実装とドキュメントを参照することで最もよく理解できます。
一般に、カーネルはテンソルよりも低いレベルで動作し、代わりにメモリに直接読み書きし、最終的には tfjs-core によってテンソルにラップされます。
カーネルが実装されたら、tfjs-core のregisterKernel関数を使用して TensorFlow.js に登録できます。カーネルを動作させたいバックエンドごとにカーネルを登録できます。登録すると、カーネルはtf.engine().runKernel(...)で呼び出すことができ、TensorFlow.js が必ず実装にディスパッチします。現在アクティブなバックエンド。
カスタムグラデーションの実装
グラデーションは通常、特定のカーネル ( tf.engine().runKernel(...)の呼び出しで使用されるのと同じカーネル名で識別されます) に対して定義されます。これにより、tfjs-core はレジストリを使用して実行時にカーネルのグラデーション定義を検索できるようになります。
カスタム グラデーションの実装は次の場合に役立ちます。
- ライブラリに存在しない可能性があるグラデーション定義の追加
- 既存の勾配定義をオーバーライドして、特定のカーネルの勾配計算をカスタマイズします。
グラデーションの実装例はここで見ることができます。
特定の呼び出しのグラデーションを実装したら、tfjs-core のregisterGradient関数を使用して TensorFlow.js に登録できます。
勾配レジストリをバイパスする (つまり、任意の関数の勾配を任意の方法で計算できるようにする) カスタム勾配を実装するもう 1 つのアプローチは、 tf.customGradを使用することです。
以下は、customGrad を使用したライブラリ内の操作の例です。