简介

本文介绍了如何将基于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

image.png

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的类型参考这里：

image.png

另外需要注意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);

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)
  ...