简介
本文介绍了如何将基于C++的分词工具cppjieba封装成一个Tensorflow Custom Op。
- 仓库地址
- 原始的cppjieba仓库
- 我们修改的分支cppjieba仓库
- 我们在深度学习开源平台DELTA中也使用了这个OP。
安装
目前支持python3.6；支持MacOS和Linux；在tf1.14下测试通过。
pip install tf-jieba
快速使用
创建一个测试目录，在该目录下：
wget https://github.com/applenob/tf_jieba/raw/master/custom_test.py
python3 custom_test.py
Tensorflow Custom Op
- 官方文档：https://www.tensorflow.org/guide/create_op
- 官方仓库：https://github.com/tensorflow/custom-op
- 可借鉴的op：https://github.com/tensorflow/lingvo/tree/master/lingvo/core/ops
编写Custom Op (Operator)就是将一些用C/C++写的特定的计算操作，封装成tensorflow平台的算子。用户可以通过不同的语言使用这些op，比如C++或者Python，就像使用官方的op一样。
工业届使用Tensorflow，一般是用python开发，用C++上线。这样，一些非Tensorflow的操作，就需要维护两个语言的版本。比如分词，Jieba是大家熟知的常用的分词工具，它的实现有python也有C++，甚至还有其他语言。但不同的实现往往具体计算会有微小差异，这样就会导致线上和线下的模型结果差异。使用OP封装就可以解决这个问题。
另外一个解决方案是将C++的代码用swig这类工具添加python接口。但是使用OP 封装，还可以将这些操作序列化（Protobufer），在所有支持tensorflow的地方都可以跑这些操作。想象一下，你把所有的数据预处理都写成op，你拿着一个SavedModel，部署到tf-serving上后，不需要其他额外代码，就可以拿到模型结果。
大致流程
目前Tensorflow官网的介绍其实已经非常详细了，建议详细阅读。我根据自己的理解再做一些补充。
编译一个自定义op主要流程是下面五步：
- 1.在c++文件中注册新的op。
- 2.用c++实现这个op。(kernel)
- 3.新建一个python的wrapper（可选）。
- 4.写一个计算该op的梯度的方式（可选）。
- 5.测试该op。
1.注册op
注册op不仅是让系统知道这个op的存在，还定义了这个op的C++版接口。
核心代码：
REGISTER_OP("JiebaCut")
.Input("sentence_in: string")
.Output("sentence_out: string")
.Attr("use_file: bool = false")
.Attr("dict_lines: list(string) = ['']")
.Attr("model_lines: list(string) = ['']")
.Attr("user_dict_lines: list(string) = ['']")
.Attr("idf_lines: list(string) = ['']")
.Attr("stop_word_lines: list(string) = ['']")
.Attr("dict_path: string = ''")
.Attr("hmm_path: string = ''")
.Attr("user_dict_path: string = ''")
.Attr("idf_path: string = ''")
.Attr("stop_word_path: string = ''")
.Attr("hmm: bool = true")
.SetShapeFn(UnchangedShape)
.Doc(R"doc(
Cut the Chines sentence into words.
sentence_in: A scalar or list of strings.
sentence_out: A scalar or list of strings.
)doc");
Input和Output是这个op的输入和输出，Attr指的是其他参数。这里注意的是输入输出和Attr的类型表示不一样。输入输出的类型更像是python中tensorflow的类型；而Attr的类型参考这里：
另外需要注意ShapeFn，用于支持tensorflow中shape inference的功能，主要实现两个功能：1.确保输入shape没有问题；2.确定输出的shape。这里使用的是UnchangedShape函数，因为输入和输出的shape是一样的。
2.实现kernel
自定义op类要继承OpKernel类。
构造函数中初始化相关参数，在构图的时候调用；compute函数中定义前向计算流程，在session run的时候调用。
void Compute(OpKernelContext* ctx) override {
std::vector<string> words;
const Tensor* t_sentence_in;
OP_REQUIRES_OK(ctx, ctx->input("sentence_in", &t_sentence_in));
Tensor* t_sentence_out;
OP_REQUIRES_OK(ctx,
ctx->allocate_output("sentence_out", t_sentence_in->shape(),
&t_sentence_out));
if (t_sentence_in->dims() == 0) {
jieba_->Cut(t_sentence_in->scalar<string>()(), words, true);
t_sentence_out->scalar<string>()() = str_util::Join(words, " ");
} else {
OP_REQUIRES(
ctx, t_sentence_in->dims() == 1,
errors::InvalidArgument("Input must be a scalar or 1D tensor."));
int batch = t_sentence_in->dim_size(0);
for (int i = 0; i < batch; i++) {
jieba_->Cut(t_sentence_in->vec<string>()(i), words, true);
t_sentence_out->vec<string>()(i) = str_util::Join(words, " ");
}
}
}
compute函数大致就是先把tensor数据结构转成C++数据结构；然后进行计算；然后再将计算结果包装成tensor数据结构返回。因此，C++数据结构和tensor的数据交换比较重要：
-
tensor到C++：tensor对象可以调用scalar<string>()或vector<string>()或matrix<string>()来获取其内部数据，然后再直接(i)或者(i,j)获取某个位置的元素的引用。 -
C++到tensor：先声明一个tensor对象，然后类似于上面的操作，将C++数据赋值给相应的引用。 - 具体操作参考这里
- 多线程和GPU相关请参考官网
实现完以后，后面还需加上注册kernel的代码。
REGISTER_KERNEL_BUILDER(Name("JiebaCut").Device(DEVICE_CPU), JiebaCutOp);
c++相关文档：https://www.tensorflow.org/api_docs/cc/class/tensorflow/tensor
3.编译
这里主要介绍使用Makefile编译的方法，使用bazel编译参考官网。
TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
g++ -std=c++11 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2
4.python封装
python封装主要实现两步：
1.将op从编译好的.so库中取出：
path = glob(os.path.join(this_directory, 'x_ops.*so'))[0]
logging.info(f'x_ops.so path: {path}')
gen_x_ops = tf.load_op_library(path)
2.设置一些参数检查，加载资源文件：
def jieba_cut(input_sentence,
use_file=True,
hmm=True):
dict_path = os.path.join(this_directory,
"./cppjieba_dict/jieba.dict.utf8")
hmm_path = os.path.join(this_directory,
"./cppjieba_dict/hmm_model.utf8")
user_dict_path = os.path.join(this_directory,
"./cppjieba_dict/user.dict.utf8")
idf_path = os.path.join(this_directory,
"./cppjieba_dict/idf.utf8")
stop_word_path = os.path.join(this_directory,
"./cppjieba_dict/stop_words.utf8")
if use_file:
output_sentence = gen_x_ops.jieba_cut(
input_sentence,
use_file=use_file,
hmm=hmm,
dict_path=dict_path,
hmm_path=hmm_path,
user_dict_path=user_dict_path,
idf_path=idf_path,
stop_word_path=stop_word_path)
...
总结
本文介绍了如何将基于C++的分词工具cppjieba封装成一个Tensorflow Custom Op。欢迎使用tf-jieba和DELTA
