tensorflow c++算子快速实践

背景

在c++算子优化的过程中，开发完的算子需要怼进TensorFlow的一套流程中(c++端开发算子，c++端和python对对接，python端封装等)。上述流程比较复杂。

用以下办法，可以快速实践c++算子的效果。

编写c++算子

按照TensorFlow c++ kernels的规则编写算子。

将c++算子打包成.so文件

TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )

g++ -std=c++14 -shared -march=native custom_op.cc -o custom_op.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

python端调用

cos = tf.load_op_library('./custom_op.so')

# 下面两句话等价
r1 = tf.split(...)
r2 = cos.custom_split(...)