瞭解Op註冊和TensorFlow中的內核鏈接

Question

我對TensorFlow相當陌生，現在正在研究自定義op開發。我已經閱讀過官方教程，但是我感覺幕後發生了很多事情，我並不總是希望將我的自定義操作放在user_ops目錄中。

因此，我拿起一個example word2vec

它使用定製的 「Skipgram」 的運算，它的註冊這裏定義：
/word2vec_o​​ps.cc
，其內核的實現是在這裏：
/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++代碼。 註冊一個op使用REGISTER_OP() macro，註冊一個內核使用REGISTER_KERNEL_BUILDER() macro。這些宏創建靜態初始化器，當包含它們的模塊被加載時運行。有兩種主要的op和kernel註冊機制：

1. 靜態鏈接到核心TensorFlow庫中，並進行靜態初始化。

2. 在運行時動態鏈接，使用tf.load_op_library()函數。

在"Skipgram"的情況下，我們使用選項1（靜態鏈接）。這些操作被鏈接到核心TensorFlow庫here中，內核鏈接在here中。 （請注意，這並不理想：word2vec操作是在我們有tf.load_op_library()之前創建的，因此沒有動態鏈接它們的機制。）因此，當您首次加載TensorFlow時（在import tensorflow as tf中），ops和內核會被註冊。如果它們是今天創建的，那麼它們將被動態加載，以便只有在需要時纔會進行註冊。 （該SyntaxNet碼具有動態加載的example。）

在巴澤勒的tf_op_gen_wrapper_py() rule需要的運算 -library依賴性的列表，並生成Python包裝那些OPS。這個規則只依賴於op註冊的原因是Python包裝完全由op註冊的操作界面決定，這是在op註冊中定義的。值得注意的是，Python接口不知道是否有特定類型或設備的專用內核。包裝生成器將op註冊鏈接到爲每個註冊操作生成Python代碼的simple C++ binary。 請注意，如果您使用tf.load_op_library()，則不需要自行調用包裝器生成器，因爲tf.load_op_library()將在運行時生成必需的代碼。