Abh*_*bhi 7 c++ swig kernel registration tensorflow
我是TensorFlow的新手,现在正在研究定制的op开发.我已经阅读了官方教程,但我觉得很多事情都发生在幕后,我并不总是想把我的自定义操作放在user_ops目录中.
因此,我举了一个例子word2vec
它使用自定义的"Skipgram"操作,其注册定义如下:
/word2vec_ops.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操作注册的幕后规则和整个编译/链接过程.
谢谢.
mrr*_*rry 20
当增加一个新的操作TensorFlow的,主要有两个步骤:
注册"op",包括定义操作的接口,以及
注册一个或多个"内核",其中涉及定义操作的实现,可能使用针对不同数据类型或设备类型(如CPU或GPU)的专门实现.
这两个步骤都涉及编写C++代码.注册op使用REGISTER_OP()宏,注册内核使用REGISTER_KERNEL_BUILDER()宏.这些宏创建静态初始化程序,这些初始化程序在加载包含它们的模块时运行.op和内核注册有两种主要的机制:
静态链接到核心TensorFlow库,并进行静态初始化.
使用该tf.load_op_library()函数在运行时动态链接.
在这种情况下"Skipgram",我们使用选项1(静态链接).在OPS被链接到核心TensorFlow图书馆在这里,和仁在链接这里.(请注意,这并不理想:word2vecops是在我们创建之前创建的tf.load_op_library(),因此没有动态链接它们的机制.)因此,当您第一次加载TensorFlow(in import tensorflow as tf)时,会注册操作和内核.如果它们是今天创建的,它们将被动态加载,这样它们只有在需要时才会被注册.(SyntaxNet代码有一个动态加载的例子.)
Bazel中的tf_op_gen_wrapper_py()规则获取了一个op -library依赖项列表,并为这些操作生成Python包装器.此规则仅依赖于op注册的原因是Python包装器完全由op的接口决定,该接口在op注册中定义.值得注意的是,Python接口不知道是否有特定类型或设备的专用内核.包装器生成器将op注册链接到一个简单的C++二进制文件,该二进制文件为每个注册的操作生成Python代码.请注意,如果使用tf.load_op_library(),则不需要自己调用包装器生成器,因为tf.load_op_library()它将在运行时生成必要的代码.