使用Tensorflow C++ API自定義操作

Tensorflow提供了大量的基本操作使得我們能夠任意組合來實現我們需要的複雜操作，但有時候我們需要的操作不太容易通過這些基本操作來組合，或者複雜的組合方式帶來嚴重的性能開銷，這時我們可能會考慮去使用py_func來包裝Python函數藉助Numpy來實現，但性能方面可能也無法達到滿意的程度，更不要說有些操作不適合向量化的寫法，這個時候用C++ API來實現自己的一個操作可能會是更好的選擇。

總得來說，我們可以通過靜態鏈接或動態連接的方式來添加我們自定義的操作，前者需要你能成功從源碼編譯安裝Tensorflow（Installing TensorFlow from Sources），後者只需要安裝Tensorflow的Python包即可，不過後者有一個缺點就是可用的C++ API是有限的，如果你需要一些別的介面可能要自己去找相關的頭文件然後手動添加進來，如果頭文件數目很多這就是一件很頭疼的事情了，但優點也是很明顯，一是簡單，二是不用強制要求使用你操作的人也要從源碼編譯安裝Tensorflow。下面主要以動態鏈接的添加為例進行敘述，對靜態鏈接添加的方式也會有涉及，二者基本上是一樣的。另外，本文主要注重實踐部分，如果你需要一些更詳細的標準化說明，可以去參考官方的Tutorial：Adding a New Op，一些比如Ops註冊的屬性說明以及如何保持Ops屬性的後向兼容性等會更加詳細。

另外如果你對Tensorflow註冊Ops和Kernel的過程或者它們是如何被Python代碼調用的感興趣，可以閱讀我之前寫的一篇文章Tensorflow是如何註冊和調用C++ New Op的。如果你想學習一下官方Ops是怎麼寫的，我推薦從Bias_Op來開始學習，因為這個Op寫的還是很清晰然後對初學者不會有太多的阻礙，當然也可以參考我最近寫的PSROIAlign。這裡還必須要推薦一個利器SourceInsight簡直是閱讀這種開源代碼的法寶。

先總結官方給的自定義op的標準流程：註冊Op，實現Op，創建python介面，實現Op梯度計算（如果不需要求導也可以直接pass掉，實現可以在python端也可以用py_func去包裝其他python函數，也可以再寫一個C++ Op來專門計算梯度），測試。

• 註冊Op

註冊Op相當於是一個聲明的過程。Op是tensorflow中非常重要的概念，一個Op接收一個或多個輸入張量，然後經過某種運算，產生其他零個或多個tensor，然後這些tensor又可以被其他Op使用。類似於C++中我們定義變數需要知道數據類型，位元組數等信息一樣，創建一個Op同樣需要一些額外信息包括attributes（輸入輸出類型以及合法取值等，也可以看作是Op的輸入但是不同於輸入的是屬性永遠是常量，其值在Op被添加到圖中時被設置，並且是一直放在CPU上的）以及輸入輸出列表，還可以直接加Doc，具體信息是我們在REGISTER_OP時指定的，REGISTER_OP是一個宏，其內部實現是一個wrapper利用了C++中的常用伎倆chaining調用實現，所有你在這添加的信息都會以另一種形式出現在動態生成的Python代碼中。有一點需要注意，在C++這邊Ops的名稱必須是CamelCase類型的，在Python那一邊會自動被轉換成Python風格的snake_case類型。

註冊這個地方還有一個SetShapeFn需要說一下，主要作用是檢查輸入的shape並指定輸出的shape，當然你也可以在Op的compute裡面檢查之類，但是這個ShapeFn有一個點是可以讓tesorflow不用執行操作就能獲取輸入輸出信息。在ShapeFn裡面你可以拿到輸入輸出的每一個維度的大小（DimensionHandle），或者屬性常量的值或者輸入常量的值，然後組合成輸出的ShapeHandle，最後調用set_output指定對應輸出的shape，同時DimensionHandle是可以做四則運算的。一開始我對如何指定輸出大小的api也有一些困惑，因為還涉及動態shape，這裡推薦仔細閱讀InferenceContext這個類，還是推薦用SourceInsight用好搜索和bookmark即可。

• 實現Op

動手之前了解一下C++中的functor、模板及其特化還是很有必要的，對lambda也有了解的話就更好了，如果你對C++不熟的話建議盡量避免使用Eigen，直接把數組取出來用C計算就行，因為tensorflow裡面的張量都是按行主序存的（多維的情況就是最外面的那一維變化最快）。

用C++實現Op有一個固定的套路，遵循這個套路可以避免走彎路，當然這都不是必須的，只要你定義了計算函數並且在kernel的Compute裡面調用你的計算即可。初學者可以參考下面這個框架來做：

1. 定義一個Functor模板類做實際的計算工作
2. 針對不同設備甚至不同數據類型特化Functor模板類
3. 定義一個Kernel模板類，繼承自OpKernel，在構造函數中根據傳進來的OpKernelConstruction設置必要的成員
4. 視情況針對別的設備特化Kernel模板類，一般無需特化，因為這個類裡面一般會做一些通用工作，然後將實際的計算轉到Functor模板類中
5. 重寫Kernel的Compute方法，利用OpKernelContext獲取輸出，分配輸出，並進行合法性檢查，然後轉調對應的計算Functor
6. 註冊Kernel

很多重要的api都在OpKernelConstruction和OpKernelContext兩個類裡面建議詳細閱讀。同時Compute方法必須是線程安全的，因此任何對類成員的訪問必須要用互斥保護。

在轉到實際的計算函數前通常會把輸入輸出Tensor的緩衝區取出來，要麼變成Eigen的表示即TensorMap（其對應的很多成員要去TensorBase裡面去找），要麼更進一步直接再調用TensorMap的data方法把緩衝區指針取出來傳給計算函數。具體地，可以去看Tensor這個類提供的一些介面。

計算的實現過程不細說，可以直接上C，可以用std::thread保證移植性，也可以利用強大的Eigen，裡面有很好的並行化機制根據每個執行單元的cost來分配線程資源，還可以利用Tensorflow提供的對Eigen進行包裝後的Shard（這個頭文件需要自己加）工具類，下面的代碼是一個使用示例：

auto work_routine = [&your_capture](int64_t start, int64_t limit){ for (int64_t worker_index = start; worker_index < limit; ++worker_index){ // do something }};const DeviceBase::CpuWorkerThreads& worker_threads = *(context->device()->tensorflow_cpu_worker_threads());const int64_t shard_cost = 4 * num_rois;Shard(worker_threads.num_threads, worker_threads.workers, total_elems, shard_cost, work_routine);

GPU的實現多數簡單情況下GetCudaLaunchConfig + CUDA_1D_KERNEL_LOOP就可以搞定，更多有用的介面可以去看cuda_kernel_helper.h，也可以藉助cub的api或者更暴力一點直接寫。順便說一下，使用Eigen或者Tensorflow的CPU並行化機制編寫代碼，通常寫出的代碼整體邏輯和GPU代碼基本一致，我在寫PSROIAlign的時候從CPU代碼移植到GPU改動的地方很少。

OpKernel的註冊和Ops的註冊比較類似，也是調用一個宏REGISTER_KERNEL_BUILDER，指定名稱、Kernel對應的設備類型，以及創建這個Kernel的C++類。詳見Tensorflow是如何註冊和調用C++ New Op的。

• 創建python介面

在這之前當然先要對自己寫的Op進行編譯，我搗鼓了一個CMakeLists感覺很好用，可以把C++和CUDA代碼分開編譯然後一起鏈接很省事，只要正確安裝了cuda和python包頭文件都可以自己找到，推薦給大家CMakeLists.txt。

編譯成功後應該可以獲得一個動態庫文件 .so，python這邊load一下然後包裝一下就好了，把官方的示例抄過來：

import tensorflow as tfzero_out_module = tf.load_op_library(./zero_out.so)zero_out = zero_out_module.zero_out

• 實現Op梯度計算

一般就是python這邊要麼用tensorflow自帶Op進行組合，要麼再去調用另一個計算梯度的自定義Op，然後整個計算過程放在一個函數裡面，用@ops.RegisterGradient修飾一下就行了，具體可參見官方文檔。

• 測試Op

同樣官方有示例，推薦一個tf.test裡面的compute_gradient和compute_gradient_error，很好用。注意這兩個函數計算的並不是梯度，而是輸出對輸入的Jacobian，計算梯度值要用tf.gradients，tf.test中的那兩個函數也是基於tf.gradients。其中compute_gradient來計算理論和數值Jacobian，compute_gradient_error計算二者之間的誤差。

• 總結

最後總結一下我目前遇到的坑：

1. 最好不要用Eigen介面，除非你對Eigen比較熟，否則建議使用前閱讀以下Lazy Evaluation and Aliasing和Eigen: Common pitfalls
2. 用Shard工具類寫CPU端的kernel，方便移植到GPU上，但要注意線程間的同步
3. 全局一致的常量輸入使用op的Attr來指定，尤其是需要基於這些常量輸入做進一步地運算的時候，因為如果將常量輸入作為Scalar類型的Tensor輸入，那麼在CPU上和GPU上運行時這些輸入將在不同的內存里，如果要基於這些常量做進一步地運算在GPU上要用cuda kernel，不利於代碼結構的簡化
4. 輸出記得先清零

此外，如果你希望靜態鏈接你的Op，那麼可以把代碼放在Tensorflow源代碼的user_ops目錄下，然後在BUILD文件里添加一個tf_custom_op_library就可以了。

暫時想到的就這麼多，最後祝大家煉丹順利~~找到理想工作。