前言
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++文件里定义的是大写