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

Question

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

因此，我拿了example word2vec

使用自定义“Skipgram”op，其注册定义如下：
的 /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

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

毕竟，如果我使用

构建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"，涉及为操作定义界面，和

2. Registering one or more "kernels"，它涉及定义操作的实现，可能是针对不同数据类型或设备类型（如CPU或GPU）的专门实现。

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

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

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

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

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