在tensorflow中使用增量編譯添加OP（不必編譯tensorflow）

關於在tensorflow(下文稱tf)中添加op，tf中文社區中有相應教程：http://www.tensorfly.cn/tfdoc/how_tos/adding_an_op.html 但該教程未提及如何將編寫的op加入tensorflow框架。
親測不必編譯tensorflow源碼，直接通過增量編譯C++和CUDA程序可以將op註冊在tensorflow中。本文以用CPU實現op爲例，記錄在tensorflow編寫、註冊和調用op的過程。

總體過程可分爲以下三步：

1. 編寫cpp文件（若使用GPU，則還需要編寫cu文件）。
2. 使用g++編譯cpp文件，生成*.so文件（若使用GPU，需要在編寫cpp文件之前使用nvcc命令編譯*.cu文件，生成*.cu.o文件）。
3. 在tf中通過*.so導入op，調用op即可。

一、編寫cpp文件。
tf中註冊op的文件主要由三部分組成——聲明op的輸入輸出及其他屬性、實現op、註冊op，典型示例代碼如下：
新建文件my_add.cc，將以下代碼複製進該文件。

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

using namespace tensorflow;
//聲明op輸入輸出等屬性
REGISTER_OP("MyAdd")
    .Input("x: int32")
    .Input("y: int32")
    .Output("z: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      c->set_output(0, c->input(1));
      return Status::OK();
    });

//實現op
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class MyAddOp : public OpKernel {
 public:
  explicit MyAddOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& a = context->input(0);
    const Tensor& b = context->input(1);
    auto A = a.flat<int32>();
    auto B = b.flat<int32>();
    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, a.shape(),
                                                     &output_tensor));
    auto output_flat = output_tensor->flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = A.size();

    for (int i = 1; i < N; i++) {
      output_flat(i) = A(i)+B(i);
    }
    output_flat(0) = 0;
  }
};

//註冊op
REGISTER_KERNEL_BUILDER(Name("MyAdd").Device(DEVICE_CPU), MyAddOp);

二、使用g++編譯，生成so文件。
主機的g++版本爲5.4.0，使用以下命令編譯該op

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++11 -shared my_add.cc -o my_add.so -fPIC -D_GLIBCXX_USE_CXX11_ABI=0 ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

若主機的g++版本號小於5，可將最後一行命令改爲

g++ -std=c++11 -shared my_add.cc -o my_add.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

執行以上命令後，在當前目錄下產生了my_add.so文件。

三、將op文件導入，調用之前編寫的cpp文件。
使用tf.load_op_library函數載入so文件，即可調用原先編寫的op。

import tensorflow as tf
so_file = 'your_add_so_file_path/my_add.so'

if __name__ == "__main__":
  #tf.test.main()
  my_add_module = tf.load_op_library(so_file) 
  out = my_add_module.my_add([5, 4, 3, 2, 1],[1, 2, 3, 4, 5])
  sess = tf.Session()
  result = sess.run(out)
  print(result)
  #output [0, 6, 6, 6, 6]

注意：註冊時tf會自動生成python_wrapper，因此在tf的python環境中調用原先op的函數名或許與之前在cpp文件中定義的函數名有所不同。在python中可以通過dir(package_name)來列出該模塊中所有函數的方法查看該op的python_wrapper。