理解 TensorFlow 中的操作注册和内核链接

Understand Op Registration and Kernel Linking in TensorFlow(理解 TensorFlow 中的操作注册和内核链接)

问题描述

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

I am fairly new to TensorFlow and right now looking into the custom op development. I have already read the official tutorial but I feel a lot of things happen behind the scenes and I do not always want to put my custom ops in user_ops directory.

因此，我使用了 示例 word2vec

As such, I took up an example word2vec

使用自定义Skipgram"操作，其注册定义如下:
/word2vec_o​​ps.cc
其内核实现在这里:
/word2vec_kernels.cc

which uses a custom "Skipgram" op whose registration is defined here:
/word2vec_ops.cc
and whose kernel implementation is here:
/word2vec_kernels.cc

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

Looking at the build file, I tried to build individual targets

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

1) bazel build -c opt tensorflow/models/embedding:word2vec_ops
This generates bunch of object files as expected.

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_pyhttps://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorflow.bzl#L197-L231

This last build uses a custom rule namely tf_op_gen_wrapper_py https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorflow.bzl#L197-L231

有趣的是，这仅取决于操作注册，而不取决于内核本身.

Interesting to note that this only depends on Op Registration and not on the kernel itself.

毕竟以上，如果我使用

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

它工作正常，但我看不到内核 C++ 代码链接的位置和方式?

it works fine, but I fail to see where and how the kernel c++ code linked?

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

Additionally, I would also like to understand the tf_op_gen_wrapper_py rule and the whole compilation/linking procedure that goes behind the scenes for ops registration.

谢谢.

推荐答案

When 为 TensorFlow 添加一种新的操作，主要有两个步骤:

When adding a new kind of operation to TensorFlow, there are two main steps:

1. 注册op"，包括为操作定义一个接口，和

1. Registering the "op", which involves defining an interface for the operation, and

注册一个或多个内核"，这涉及为操作定义实现，可能具有针对不同数据类型或设备类型(如 CPU 或 GPU)的专门实现.

Registering one or more "kernels", which involves defining implementation(s) for the operation, perhaps with specialized implementations for different data types, or device types (like CPU or GPU).

这两个步骤都涉及编写 C++ 代码.注册操作使用 REGISTER_OP 宏，注册内核使用 REGISTER_KERNEL_BUILDER() 宏.这些宏创建静态初始化程序，在加载包含它们的模块时运行.op和kernel注册主要有两种机制:

Both steps involve writing C++ code. Registering an op uses the REGISTER_OP() macro, and registering a kernel uses the REGISTER_KERNEL_BUILDER() macro. These macros create static initializers that run when the module containing them is loaded. There are two main mechanisms for op and kernel registration:

1. 静态链接到核心 TensorFlow 库，以及静态初始化.

1. Static linking into the core TensorFlow library, and static initialization.

运行时动态链接，使用 tf.load_op_library() 函数.

Dynamic linking at runtime, using the tf.load_op_library() function.

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

In the case of "Skipgram", we use option 1 (static linking). The ops are linked into the core TensorFlow library here, and the kernels are linked in here. (Note that this is not ideal: the word2vec ops were created before we had tf.load_op_library(), and so there was no mechanism for linking them dynamically.) Hence the ops and kernels are registered when you first load TensorFlow (in import tensorflow as tf). If they were created today, they would be dynamically loaded, such that they would only be registered if they were needed. (The SyntaxNet code has an example of dynamic loading.)

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