如何将分词工具Jieba封装成一个Tensorflow Operator

Posted on | In | Visitors:

简介

本文介绍了如何将基于C++的分词工具cppjieba封装成一个Tensorflow Custom Op。

安装

目前支持python3.6；支持MacOS和Linux；在tf1.14下测试通过。

1
pip install tf-jieba

快速使用

创建一个测试目录，在该目录下：

1
2
wget https://github.com/applenob/tf_jieba/raw/master/custom_test.py
python3 custom_test.py

image

Tensorflow Custom Op

编写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++版接口。

核心代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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

另外需要注意ShapeFn，用于支持tensorflow中shape inference的功能，主要实现两个功能：1.确保输入shape没有问题；2.确定输出的shape。这里使用的是UnchangedShape函数，因为输入和输出的shape是一样的。

2.实现kernel

自定义op类要继承OpKernel类。

构造函数中初始化相关参数，在构图的时候调用；compute函数中定义前向计算流程，在session run的时候调用。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
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的代码。

1
REGISTER_KERNEL_BUILDER(Name("JiebaCut").Device(DEVICE_CPU), JiebaCutOp);

c++相关文档

3.编译

这里主要介绍使用Makefile编译的方法，使用bazel编译参考官网。

1
2
3
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库中取出：

1
2
3
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.设置一些参数检查：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def jieba_cut(input_sentence,
              use_file=True,
              dict_path="",
              hmm_path="",
              user_dict_path="",
              idf_path="",
              stop_word_path="",
              dict_lines=[""],
              model_lines=[""],
              user_dict_lines=[""],
              idf_lines=[""],
              stop_word_lines=[""],
              hmm=True):
  if use_file:
    if not os.path.exists(dict_path):
      raise FileNotFoundError("Dict file not found: {}!".format(dict_path))

    if not os.path.exists(hmm_path):
      raise FileNotFoundError("HMM Model file not found: {}!".format(hmm_path))

    if not os.path.exists(user_dict_path):
      raise FileNotFoundError("User dict file not found: {}!".format(user_dict_path))

    if not os.path.exists(idf_path):
      raise FileNotFoundError("IDF file not found: {}!".format(idf_path))

    if not os.path.exists(stop_word_path):
      raise FileNotFoundError("Stop words file not found: {}!".format(stop_word_path))
  ...

总结

本文介绍了如何将基于C++的分词工具cppjieba封装成一个Tensorflow Custom Op。欢迎使用tf-jieba和DELTA