TensorFlow实现自定义Op方式

2023-08-02 09:46:04 78

『写在前面』

以CTCBeamsearchdecoder为例，简单整理一下TensorFlow实现自定义Op的操作流程。

基本的流程

1.定义Op接口

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

REGISTER_OP("Custom")
.Input("custom_input:int32")
.Output("custom_output:int32");

2.为Op实现Compute操作(CPU)或实现kernel(GPU)

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

usingnamespacetensorflow;

classCustomOp:publicOpKernel{
public:
explicitCustomOp(OpKernelConstruction*context):OpKernel(context){}
voidCompute(OpKernelContext*context)override{
//获取输入tensor.
constTensor&input_tensor=context->input(0);
autoinput=input_tensor.flat();
//创建一个输出tensor.
Tensor*output_tensor=NULL;
OP_REQUIRES_OK(context,context->allocate_output(0,input_tensor.shape(),
&output_tensor));
autooutput=output_tensor->templateflat();
//进行具体的运算，操作input和output
//……
}
};

3.将实现的kernel注册到TensorFlow系统中

REGISTER_KERNEL_BUILDER(Name("Custom").Device(DEVICE_CPU),CustomOp);

CTCBeamSearchDecoder自定义

该Op对应TensorFlow中的源码部分

Op接口的定义：

tensorflow-master/tensorflow/core/ops/ctc_ops.cc

CTCBeamSearchDecoder本身的定义：

tensorflow-master/tensorflow/core/util/ctc/ctc_beam_search.cc

Op-Class的封装与Op注册：

tensorflow-master/tensorflow/core/kernels/ctc_decoder_ops.cc

基于源码修改的Op

#include
#include
#include

#include"tensorflow/core/util/ctc/ctc_beam_search.h"

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

namespacetf=tensorflow;
usingtf::shape_inference::DimensionHandle;
usingtf::shape_inference::InferenceContext;
usingtf::shape_inference::ShapeHandle;

usingnamespacetensorflow;

REGISTER_OP("CTCBeamSearchDecoderWithParam")
.Input("inputs:float")
.Input("sequence_length:int32")
.Attr("beam_width:int>=1")
.Attr("top_paths:int>=1")
.Attr("merge_repeated:bool=true")
//新添加了两个参数
.Attr("label_selection_size:int>=0=0")
.Attr("label_selection_margin:float")
.Output("decoded_indices:top_paths*int64")
.Output("decoded_values:top_paths*int64")
.Output("decoded_shape:top_paths*int64")
.Output("log_probability:float")
.SetShapeFn([](InferenceContext*c){
ShapeHandleinputs;
ShapeHandlesequence_length;

TF_RETURN_IF_ERROR(c->WithRank(c->input(0),3,&inputs));
TF_RETURN_IF_ERROR(c->WithRank(c->input(1),1,&sequence_length));

//Getbatchsizefrominputsandsequence_length.
DimensionHandlebatch_size;
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(inputs,1),c->Dim(sequence_length,0),&batch_size));

int32top_paths;
TF_RETURN_IF_ERROR(c->GetAttr("top_paths",&top_paths));

//Outputs.
intout_idx=0;
for(inti=0;iset_output(out_idx++,c->Matrix(InferenceContext::kUnknownDim,2));
}
for(inti=0;iset_output(out_idx++,c->Vector(InferenceContext::kUnknownDim));
}
ShapeHandleshape_v=c->Vector(2);
for(inti=0;iset_output(out_idx++,shape_v);
}
c->set_output(out_idx++,c->Matrix(batch_size,top_paths));
returnStatus::OK();
});

typedefEigen::ThreadPoolDeviceCPUDevice;

inlinefloatRowMax(constTTypes::UnalignedConstMatrix&m,intr,
int*c){
*c=0;
CHECK_LT(0,m.dimension(1));
floatp=m(r,0);
for(inti=1;ip){
p=m(r,i);
*c=i;
}
}
returnp;
}

classCTCDecodeHelper{
public:
CTCDecodeHelper():top_paths_(1){}

inlineintGetTopPaths()const{returntop_paths_;}
voidSetTopPaths(inttp){top_paths_=tp;}

StatusValidateInputsGenerateOutputs(
OpKernelContext*ctx,constTensor**inputs,constTensor**seq_len,
Tensor**log_prob,OpOutputList*decoded_indices,
OpOutputList*decoded_values,OpOutputList*decoded_shape)const{
Statusstatus=ctx->input("inputs",inputs);
if(!status.ok())returnstatus;
status=ctx->input("sequence_length",seq_len);
if(!status.ok())returnstatus;

constTensorShape&inputs_shape=(*inputs)->shape();

if(inputs_shape.dims()!=3){
returnerrors::InvalidArgument("inputsisnota3-Tensor");
}

constint64max_time=inputs_shape.dim_size(0);
constint64batch_size=inputs_shape.dim_size(1);

if(max_time==0){
returnerrors::InvalidArgument("max_timeis0");
}
if(!TensorShapeUtils::IsVector((*seq_len)->shape())){
returnerrors::InvalidArgument("sequence_lengthisnotavector");
}

if(!(batch_size==(*seq_len)->dim_size(0))){
returnerrors::FailedPrecondition(
"len(sequence_length)!=batch_size.","len(sequence_length):",
(*seq_len)->dim_size(0),"batch_size:",batch_size);
}

autoseq_len_t=(*seq_len)->vec();

for(intb=0;boutput_list("decoded_indices",decoded_indices);
if(!s.ok())returns;
s=ctx->output_list("decoded_values",decoded_values);
if(!s.ok())returns;
s=ctx->output_list("decoded_shape",decoded_shape);
if(!s.ok())returns;

returnStatus::OK();
}

//sequences[b][p][ix]storesdecodedvalue"ix"ofpath"p"forbatch"b".
StatusStoreAllDecodedSequences(
conststd::vector>>&sequences,
OpOutputList*decoded_indices,OpOutputList*decoded_values,
OpOutputList*decoded_shape)const{
//Calculatethetotalnumberofentriesforeachpath
constint64batch_size=sequences.size();
std::vectornum_entries(top_paths_,0);

//Calculatenum_entriesperpath
for(constauto&batch_s:sequences){
CHECK_EQ(batch_s.size(),top_paths_);
for(intp=0;pallocate(p,TensorShape({p_num,2}),&p_indices);
if(!s.ok())returns;
s=decoded_values->allocate(p,TensorShape({p_num}),&p_values);
if(!s.ok())returns;
s=decoded_shape->allocate(p,TensorShape({2}),&p_shape);
if(!s.ok())returns;

autoindices_t=p_indices->matrix();
autovalues_t=p_values->vec();
autoshape_t=p_shape->vec();

int64max_decoded=0;
int64offset=0;

for(int64b=0;bGetAttr("merge_repeated",&merge_repeated_));
OP_REQUIRES_OK(ctx,ctx->GetAttr("beam_width",&beam_width_));
//从参数列表中读取新添的两个参数
OP_REQUIRES_OK(ctx,ctx->GetAttr("label_selection_size",&label_selection_size));
OP_REQUIRES_OK(ctx,ctx->GetAttr("label_selection_margin",&label_selection_margin));
inttop_paths;
OP_REQUIRES_OK(ctx,ctx->GetAttr("top_paths",&top_paths));
decode_helper_.SetTopPaths(top_paths);
}

voidCompute(OpKernelContext*ctx)override{
constTensor*inputs;
constTensor*seq_len;
Tensor*log_prob=nullptr;
OpOutputListdecoded_indices;
OpOutputListdecoded_values;
OpOutputListdecoded_shape;
OP_REQUIRES_OK(ctx,decode_helper_.ValidateInputsGenerateOutputs(
ctx,&inputs,&seq_len,&log_prob,&decoded_indices,
&decoded_values,&decoded_shape));

autoinputs_t=inputs->tensor();
autoseq_len_t=seq_len->vec();
autolog_prob_t=log_prob->matrix();

constTensorShape&inputs_shape=inputs->shape();

constint64max_time=inputs_shape.dim_size(0);
constint64batch_size=inputs_shape.dim_size(1);
constint64num_classes_raw=inputs_shape.dim_size(2);
OP_REQUIRES(
ctx,FastBoundsCheck(num_classes_raw,std::numeric_limits::max()),
errors::InvalidArgument("num_classescannotexceedmaxint"));
constintnum_classes=static_cast(num_classes_raw);

log_prob_t.setZero();

std::vector::UnalignedConstMatrix>input_list_t;

for(std::size_tt=0;tbeam_search(num_classes,beam_width_,
&beam_scorer_,1/*batch_size*/,
merge_repeated_);
//使用传入的两个参数进行Set
beam_search.SetLabelSelectionParameters(label_selection_size,label_selection_margin);
Tensorinput_chip(DT_FLOAT,TensorShape({num_classes}));
autoinput_chip_t=input_chip.flat();

std::vector>>best_paths(batch_size);
std::vectorlog_probs;

//Assumption:theblankindexisnum_classes-1
for(intb=0;b(input_chip_t.data(),num_classes);
beam_search.Step(input_bi);
}
OP_REQUIRES_OK(
ctx,beam_search.TopPaths(decode_helper_.GetTopPaths(),&best_paths_b,
&log_probs,merge_repeated_));

beam_search.Reset();

for(intbp=0;bp::DefaultBeamScorerbeam_scorer_;
boolmerge_repeated_;
intbeam_width_;
//新添两个数据成员，用于存储新加的参数
intlabel_selection_size;
floatlabel_selection_margin;
TF_DISALLOW_COPY_AND_ASSIGN(CTCBeamSearchDecoderWithParamOp);
};

REGISTER_KERNEL_BUILDER(Name("CTCBeamSearchDecoderWithParam").Device(DEVICE_CPU),
CTCBeamSearchDecoderWithParamOp);

将自定义的Op编译成.so文件

在tensorflow-master目录下新建一个文件夹custom_op

cdcustom_op

新建一个BUILD文件，并在其中添加如下代码：

cc_library(
name="ctc_decoder_with_param",
srcs=[
"new_beamsearch.cc"
]+
glob(["boost_locale/**/*.hpp"]),
includes=["boost_locale"],
copts=["-std=c++11"],
deps=["//tensorflow/core:core",
"//tensorflow/core/util/ctc",
"//third_party/eigen3",
],
)

编译过程:

1.cd到tensorflow-master目录下

2.bazelbuild-copt--copt=-O3//tensorflow:libtensorflow_cc.so//custom_op:ctc_decoder_with_param

3.bazel-bin/custom_op目录下生成libctc_decoder_with_param.so

在训练（预测）程序中使用自定义的Op

在程序中定义如下的方法：

decode_param_op_module=tf.load_op_library('libctc_decoder_with_param.so')
defdecode_with_param(inputs,sequence_length,beam_width=100,
top_paths=1,merge_repeated=True):
decoded_ixs,decoded_vals,decoded_shapes,log_probabilities=(
decode_param_op_module.ctc_beam_search_decoder_with_param(
inputs,sequence_length,beam_width=beam_width,
top_paths=top_paths,merge_repeated=merge_repeated,
label_selection_size=40,label_selection_margin=0.99))
return(
[tf.SparseTensor(ix,val,shape)for(ix,val,shape)
inzip(decoded_ixs,decoded_vals,decoded_shapes)],
log_probabilities)

然后就可以像使用tf.nn.ctc_beam_search_decoder一样使用该Op了。

以上这篇TensorFlow实现自定义Op方式就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。

声明：本文内容来源于网络，版权归原作者所有，内容由互联网用户自发贡献自行上传，本网站不拥有所有权，未作人工编辑处理，也不承担相关法律责任。如果您发现有涉嫌版权的内容，欢迎发送邮件至：czq8825#qq.com（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。

TensorFlow实现自定义Op方式

热门推荐

随机推荐