c++ - 了解 TensorFlow 中的操作注册和内核链接

Question

我是 TensorFlow 的新手，现在正在研究自定义操作开发。我已经阅读了官方教程，但我觉得幕后发生了很多事情，我并不总是想将我的自定义操作放在 user_ops 目录中。

因此，我占用了一个 example word2vec

它使用自定义的“Skipgram”操作，其注册在此处定义:
/word2vec_o​​ps.cc
其内核实现在这里:
/word2vec_kernels.cc

查看构建文件，我尝试构建单个目标

1) bazel build -c opt tensorflow/models/embedding:word2vec_o​​ps
这会按预期生成一堆目标文件。

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

有趣的是，这仅取决于 Op Registration 而不是内核本身。

毕竟，如果我构建 py_binary 本身使用

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

它工作正常，但我看不到内核 C++ 代码在哪里以及如何链接？

此外，我还想了解 tf_op_gen_wrapper_py 规则以及操作注册幕后的整个编译/链接过程。

谢谢。

Answer 1

当 adding a new kind of operation to TensorFlow ，主要有两个步骤:

1. Registering the "op" ，其中涉及为操作定义接口(interface)，和

2. Registering one or more "kernels" ，这涉及为操作定义实现，可能具有针对不同数据类型或设备类型(如 CPU 或 GPU)的专门实现。

这两个步骤都涉及编写 C++ 代码。注册操作使用 REGISTER_OP() macro , 并注册内核使用 REGISTER_KERNEL_BUILDER() macro .这些宏创建静态初始值设定项，当加载包含它们的模块时运行。操作和内核注册有两种主要机制:

1. 静态链接到核心 TensorFlow 库和静态初始化。

2. 运行时动态链接，使用 tf.load_op_library()功能。

对于 "Skipgram"，我们使用选项 1(静态链接)。操作链接到核心 TensorFlow 库 here ，内核链接在 here 中. (请注意，这并不理想:word2vec 操作是在我们拥有 tf.load_op_library() 之前创建的，因此没有动态链接它们的机制。)因此操作和内核在您首次加载 TensorFlow 时注册(在 import tensorflow as tf 中)。如果它们是今天创建的，它们将被动态加载，这样它们只会在需要时才被注册。 (SyntaxNet 代码有一个 example 的动态加载。)

tf_op_gen_wrapper_py() rule在 Bazel 中获取 op 库依赖项列表，并为这些操作生成 Python 包装器。该规则仅取决于操作注册的原因是 Python 包装器完全由操作的接口(interface)决定，该接口(interface)在操作注册中定义。值得注意的是，Python 接口(interface)不知道是否有针对特定类型或设备的专用内核。包装器生成器将操作注册链接到 simple C++ binary 中。为每个注册的操作生成 Python 代码。请注意，如果您使用 tf.load_op_library()，则无需自己调用包装器生成器，因为 tf.load_op_library() 会在运行时生成必要的代码.