TensorFlowでOp登録とカーネルリンクを理解する

Question

私はTensorFlowを初めて使っていて、今すぐカスタムop開発を検討しています。私はすでに公式のチュートリアルを読んでいますが、多くのことが舞台裏で起こっていると感じています。私はいつもカスタム操作をuser_opsディレクトリに入れたいとは思っていません。
/word2vec_o​​ps.cc
とそのカーネルの実装ここにある：
など、私が登録ここで定義されたカスタム "Skipgram" のOPを使用していますexample word2vec

を取り上げたとして

/word2vec_kernels.cc

ビルドファイルを見ると、個別のターゲットを作成しようとしました。

1）bazel build -c opt tensorflow/models/embedding:word2vec_ops
これは、期待どおりにオブジェクトファイルの束を生成します。

2）bazel build -c opt tensorflow/models/embedding:word2vec_kernels
これと同じです。

3）bazel build -c opt tensorflow/models/embedding:word2vec_kernels:gen_word2vec

この最後のビルドではこれが唯一のオペアンプ登録ではなく、カーネル自体に依存することに注意することは、すなわちtf_op_gen_wrapper_py https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorflow.bzl#L197-L231

興味深いカスタムルールを使用しています。上記のすべての後

、私は

bazel build -c opt tensorflow/models/embedding:word2vec

を使用して自分自身py_binaryを構築する場合、それは正常に動作しますが、私はどこで、どのようにカーネルのC++コードをリンク参照に失敗？

さらに、tf_op_gen_wrapper_pyルールと、ops登録の背景にあるコンパイル/リンク手順全体についても理解したいと思います。

ありがとうございました。

Answer 1

adding a new kind of operation to TensorFlowは、二つの主要な段階がある：

1. Registering the "op"、操作するためのインタフェースを定義することを含む、と

2. Registering one or more "kernels"、の実装（複数可）を定義することを含みますおそらく、さまざまなデータタイプやデバイスタイプ（CPUやGPUなど）に特化した実装で動作します。

両方の手順では、C++コードを記述します。 オペレーションを登録するとREGISTER_OP() macroが使用され、カーネルを登録するにはREGISTER_KERNEL_BUILDER() macroが使用されます。これらのマクロは、それらを含むモジュールがロードされたときに実行される静的イニシャライザを作成します。opとカーネルの登録には、主に2つのメカニズムがあります。

1. コアTensorFlowライブラリへの静的リンクと静的初期化。

2. tf.load_op_library()機能を使用して、実行時のダイナミックリンク。

"Skipgram"の場合、オプション1（静的リンク）を使用します。 opsはコアTensorFlowライブラリhereにリンクされ、カーネルはhereにリンクされています。 （これは理想的ではないことに注意してください：tf.load_op_library()の前にword2vecのopsが作成されているため、動的にリンクするメカニズムはありませんでした。）したがって、最初にTensorFlowをロードしたときにopsとカーネルが登録されます（import tensorflow as tf）。今日作成されたものは、動的にロードされるため、必要な場合にのみ登録されます。 （SyntaxNetコードは、動的ローディングのexampleを有する。）

Bazelでtf_op_gen_wrapper_py() ruleはOP -library依存関係のリストを受け取り、それらのオペレーションのためのPythonラッパーを生成します。このルールがop登録のみに依存する理由は、Pythonラッパーがop登録で定義されたopのインターフェースによって完全に決まるからです。特に、Pythonインタフェースでは、特定のタイプやデバイスに特化したカーネルがあるかどうかはわかりません。ラッパージェネレータは、オペレーション登録を、登録されたオペレーションごとにPythonコードを生成するsimple C++ binaryにリンクします。 tf.load_op_library()を使用する場合は、実行時に必要なコードをtf.load_op_library()が生成するため、ラッパージェネレータを自分で呼び出す必要はありません。