前言
Tensorflow几年前已经开始用了，之前一直在数据量不大的场景用，而且没有上线serving，很多坑体会不到。最近接手新的项目，重新捡起TF，踏上了不断踩坑的旅程。
自定义OP
使用C++开发自定义op的动机是，在使用tf.dataset 对原始输入的文本数据进行处理，发现性能实在是奇慢无比。猜测可能是封装好的通用方法，实现了许多对当前使用场景冗余的逻辑，于是决定自己开发一个自定义op 来实现decode_csv的功能。
先明确一下输入和输出，这个函数我是放在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) 就可以获取到。
- 输出的tensor需要先定义Tensor，确定shape 然后通过OP_REQUIRES_OK()这个方法初始化对应形状的向量。最后输出的内容是通过Tensor对象里的tensor成员变量来定义的，这里<>里的第二个参数表示输出向量的维度，必须要和上面shape里定义维度一致。输出的结果直接写入该成员变量里即可。
- 还需通过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++文件里定义的是大写