了解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()将在运行时生成必需的代码。