前言

Tensorflow幾年前已經(jīng)開始用了，之前一直在數(shù)據(jù)量不大的場景用，而且沒有上線serving，很多坑體會不到。最近接手新的項目，重新?lián)炱餞F，踏上了不斷踩坑的旅程。

自定義OP

使用C++開發(fā)自定義op的動機是，在使用tf.dataset 對原始輸入的文本數(shù)據(jù)進行處理，發(fā)現(xiàn)性能實在是奇慢無比。猜測可能是封裝好的通用方法，實現(xiàn)了許多對當前使用場景冗余的邏輯，于是決定自己開發(fā)一個自定義op 來實現(xiàn)decode_csv的功能。

先明確一下輸入和輸出，這個函數(shù)我是放在dataset.map()中使用。

dataset使用

map函數(shù)

一開始是打算先map再batch 這樣是對每一行進行處理 ，然而發(fā)現(xiàn)這樣做之后速度還是慢，因為輸入的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 <cstring>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <vector>
#include <string>
#include <iostream>

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<string>();
        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<float,2>();
        auto label_output = label->tensor<float,2>();
        auto weight_output = weight->tensor<float,2>();

        for(int i=0;i<input_tensor1.shape().dim_size(0);i++){
        int output_idx = 0;
        string::size_type p = 0;
        int feature_num = 0;
        float weights =0.0;
        string line = input1(i);
        while(p!=line.size()) {
            if(output_idx ==0){
                    if(feature_num==1976){
                        ++output_idx;
                    }
                    else if(line[p]!=','){
                        feature_output(i,feature_num)=1.0*(line[p]-'0');
                        ++feature_num;

                    }

            }
            if(output_idx==1 && line[p]!=','){
                    ++output_idx;
                    label_output(i,0) = 1.0*(line[p]-'0');

            }
            if(output_idx==2 && line[p]!=','){
                    weights= weights*10 + 1.0*(line[p]-'0');
            }
            p++;
        }
        weight_output(i,0)=weights;
        }
    }
};


REGISTER_KERNEL_BUILDER(Name("Fextract").Device(DEVICE_CPU), FeaturesExtractOp);

}// namespace tensorflow

• 最開始要使用REGISTER_OP注冊這個op，可以在這里定義input output還有shape，attribute等。這里因為輸入沒有帶參數(shù)，所以沒有attribute，只有input，output, shape. 值得一提的是，SetShapeFn里 輸出是output的行數(shù)是通過獲取輸入shape[0]得到的，輸出都是二維的，所以可以使用Matrix.
  TF_RETURN_IF_ERROR 是對輸入格式進行檢查，這里因為對batch_size行進行處理，通過batch()函數(shù)轉(zhuǎn)成了一個1維的list,因此這里檢查行數(shù)是1。也可以不寫這句，但是為了確保使用安全，最好還是檢查一下。 shape_inference::DimensionHandle row_shape = c->Dim(input_shape, 0); 這句就是獲取shape[0]的過程，輸入?yún)?shù)0表示獲取shape的第0位。

• 接下來開始寫實現(xiàn)op的類了。按照模板，先定義一個構(gòu)造方法，因為我們沒有傳入?yún)?shù)，所以這里默認就是空的。由于是繼承了OpKernel，所以還是需要把context傳給父類。

• 真正的計算過程實現(xiàn)在Compute方法中，代碼還是比較清晰的。需要注意點是
  1 . 因為我們的輸入是batch_size行的數(shù)據(jù)，所以需要在代碼里獲取到這個信息input_tensor1.shape().dim_size(0) 就可以獲取到。

2. 輸出的tensor需要先定義Tensor，確定shape 然后通過OP_REQUIRES_OK()這個方法初始化對應形狀的向量。最后輸出的內(nèi)容是通過Tensor對象里的tensor成員變量來定義的，這里<>里的第二個參數(shù)表示輸出向量的維度，必須要和上面shape里定義維度一致。輸出的結(jié)果直接寫入該成員變量里即可。
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即可產(chǎn)生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++文件里定義的是大寫