自定义Tensorflow OP开发经验总结

前言

Tensorflow几年前已经开始用了，之前一直在数据量不大的场景用，而且没有上线serving，很多坑体会不到。最近接手新的项目，重新捡起TF，踏上了不断踩坑的旅程。

自定义OP

使用C++开发自定义op的动机是，在使用tf.dataset 对原始输入的文本数据进行处理，发现性能实在是奇慢无比。猜测可能是封装好的通用方法，实现了许多对当前使用场景冗余的逻辑，于是决定自己开发一个自定义op 来实现decode_csv的功能。

先明确一下输入和输出，这个函数我是放在dataset.map()中使用。

dataset使用

map函数

一开始是打算先map再batch 这样是对每一行进行处理 ，然而发现这样做之后速度还是慢，因为输入的txt文件太大，先对每一行处理完效率太低。所以改成先batch再map的方式。因此输入就是batch_size 行文本，输出是对应batch的feature，label和weight 。都是二维Tensor

op的写法

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"

#include 
#include 
#include 
#include 
#include 
#include 
#include 

namespace tensorflow {

REGISTER_OP("Fextract")
    .Input("line: string")
    .Output("feature: float32")
    .Output("label: float32")
    .Output("weight: float32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    shape_inference::ShapeHandle input_shape;
    TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input_shape));

    shape_inference::DimensionHandle row_shape = c->Dim(input_shape, 0);

    c->set_output(0, c->Matrix(row_shape, 1976));
    c->set_output(1, c->Matrix(row_shape, 1));
    c->set_output(2, c->Matrix(row_shape, 1));
    return Status::OK();
  });

class FeaturesExtractOp : public OpKernel {
    public:
        explicit FeaturesExtractOp(OpKernelConstruction* context) : OpKernel(context) {
        }

    void Compute(OpKernelContext* context) override {
        const Tensor& input_tensor1 = context->input(0);
        auto input1 = input_tensor1.flat();
        Tensor * feature = NULL;
        Tensor * label  = NULL;
        Tensor * weight  = NULL;

        TensorShape feature_shape({input_tensor1.shape().dim_size(0),1976});
        TensorShape label_shape({input_tensor1.shape().dim_size(0),1});
        TensorShape weight_shape({input_tensor1.shape().dim_size(0),1});

        OP_REQUIRES_OK(context, context->allocate_output(0, feature_shape, &feature));
        OP_REQUIRES_OK(context, context->allocate_output(1, label_shape, &label));
        OP_REQUIRES_OK(context, context->allocate_output(2, weight_shape, &weight));

        auto feature_output = feature->tensor();
        auto label_output = label->tensor();
        auto weight_output = weight->tensor();

        for(int i=0;i

• 最开始要使用REGISTER_OP注册这个op，可以在这里定义input output还有shape，attribute等。这里因为输入没有带参数，所以没有attribute，只有input，output, shape. 值得一提的是，SetShapeFn里 输出是output的行数是通过获取输入shape[0]得到的，输出都是二维的，所以可以使用Matrix.
  TF_RETURN_IF_ERROR 是对输入格式进行检查，这里因为对batch_size行进行处理，通过batch()函数转成了一个1维的list,因此这里检查行数是1。也可以不写这句，但是为了确保使用安全，最好还是检查一下。 shape_inference::DimensionHandle row_shape = c->Dim(input_shape, 0); 这句就是获取shape[0]的过程，输入参数0表示获取shape的第0位。

• 接下来开始写实现op的类了。按照模板，先定义一个构造方法，因为我们没有传入参数，所以这里默认就是空的。由于是继承了OpKernel，所以还是需要把context传给父类。

• 真正的计算过程实现在Compute方法中，代码还是比较清晰的。需要注意点是
  1 . 因为我们的输入是batch_size行的数据，所以需要在代码里获取到这个信息input_tensor1.shape().dim_size(0) 就可以获取到。

2. 输出的tensor需要先定义Tensor，确定shape 然后通过OP_REQUIRES_OK()这个方法初始化对应形状的向量。最后输出的内容是通过Tensor对象里的tensor成员变量来定义的，这里<>里的第二个参数表示输出向量的维度，必须要和上面shape里定义维度一致。输出的结果直接写入该成员变量里即可。
3. 还需通过REGISTER_KERNEL_BUILDER定义方法名，这个是在Python里使用的时候的名字。将方法名和上面的类进行绑定。

打包和使用

all: tfop

TF_INC=/home/recommend/.local/lib/python2.7/site-packages/tensorflow/include
TF_LIB=/home/recommend/.local/lib/python2.7/site-packages/tensorflow


tfop:
    g++ -D_GLIBCXX_USE_CXX11_ABI=0 -DEXTRACT -I. -std=c++11 -shared extracts_op.cc -o extracts_op.so -fPIC -I$(TF_INC) -I$(TF_INC)/external/nsync/public -L$(TF_LIB) -ltensorflow_framework -g

clean:
    rm -f *.o *.pyc

写一个上面的Makefile，输入是上面的源码extracts_op.cc 输出定义为so文件，还有配置好上面的tf的lib路径 输入make即可产生so文件。

import os
import tensorflow as tf
library_filename = os.path.join(tf.resource_loader.get_data_files_path(),'./extracts_op.so')
extract_op_module = tf.load_op_library(library_filename)
....
feature,label,weight = extract_op_module.fextract(line)
....

注意使用的方法名首字母要小写，之前在c++文件里定义的是大写