Tensorflow 中怎麼定義自己的層呢？

01-14

問題是這個樣子的
在tensorflow中，常規操作（加法，乘法）等，可以自動反向求導，但是問題是，如果包含一些特殊的操作的時候，自己定義一些特殊的層，tensorflow是不能自動反向求導的，那自己要怎麼做才能使的反向求導順利完成呢？
希望我有描述清楚我的問題

在Eager Execution出來之前，可以使用py_func方法自定義層，但是可惜只能在CPU上運行。如果要在GPU上運行需要用C++編寫。

在Eager Execution出來之後，可以輕鬆定義層和梯度，在定義時只要使用tf.matmul或者tf.multiply等tf的內置函數進行計算，自定義的層或者梯度就可以在GPU上執行。

可以參考Github上的TensorFlow Eager Execution示例教程：https://github.com/ZhuanZhiCode/TensorFlow-Eager-Execution-Examples

其中有一個自定義softmax操作，並為softmax自定義梯度的例子：https://github.com/ZhuanZhiCode/TensorFlow-Eager-Execution-Examples/blob/master/examples/CustomGradient/custom_softmax.py

注意，在Eager Execution機制中，你只要用tf內置操作定義前向傳播，tf就會自動會你定義的函數計算梯度，但很多情況下，自動計算梯度的效率，不如你人工化簡公示後自定義的梯度計算方法來的快。比較典型的例子就是Softmax，使用交叉熵的Softmax的梯度化簡後就是object - target，如果用自動梯度，那就比這個複雜的多了。因此在很多情況下，自定義梯度是一個很重要的需求。

自問自答： 2017.11.01 tensorlfow引入動態圖機制 eager execution 為自定義和高階梯度提供強大支持。。！！撒花，撒花，撒花

碰巧看到這個問題，前段時間做比賽順便學習了一下用C++自定義tensorflow的op，比賽雖然做的不怎麼樣，但是學了不少額外的東西。。。

今天順便總結一下，希望對初次接觸的同學有所幫助。這裡僅涉及動態鏈接的op添加，因為當時我使用的bazel有問題我也懶得降版本，導致無法從源碼進行編譯。動態添加是最簡單的方式，只要安裝了tensorflow的python包就可以編譯然後鏈接使用，有一個缺點就是可用的api是有限的。

講真的其實用C++寫一個tensorflow的新操作還是很簡單的，甚至只要會C然後利用一個好的框架就大致可以搞定，總的來說可以參考的資料主要有：tensorflow offical add new op 以及tensorflow源碼中官方的op是怎麼添加的，其實都是一種套路，多學習官方代碼里的op事半功倍，具體來說我推薦Bias_Op來開始學習，因為這個op寫的還是很清晰然後對初學者不會有太多的阻礙。這裡還必須要推薦一個利器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的輸入但是不同於輸入的是屬性永遠是常量）以及輸入輸出列表，還可以直接加Doc，具體信息是我們在REGISTER_OP時指定的，REGISTER_OP是一個宏，其內部實現是一個wrapper利用了C++中的常用伎倆chaining調用實現。具體註冊的過程可以參看我之前寫的一個簡單的介紹 -- 知乎用戶：tensorflow中的op註冊的原理是什麼？

註冊這個地方還有一個SetShapeFn需要說一下，主要作用是檢查輸入的shape並指定輸出的shape，當然你也可以在op的compute裡面檢查之類，但是這個ShapeFn有一個點是可以讓tesorflow不用執行操作就能獲取輸入輸出信息。一開始我對如何指定輸出大小的api也有一些困惑，因為還涉及動態shape，這裡推薦仔細閱讀InferenceContext這個類，還是推薦用SourceInsight用好搜索和bookmark即可。

實現op

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

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

每個Op繼承自OpKernel，然後最後在註冊kernel的時候可以對不同的設備類型進行特化（如果你CPU和GPU實現相差很多，那麼最好是定義一個對多數設備通用的類模板，然後對特殊設備特化這個模板）。OpKernel::Compute()是每個op計算的入口，每個op都要進行重寫，通常取出所有的輸入做一些常規檢查，然後分配輸出和臨時tensor，最後轉去調用實際的計算函數。同樣很多重要的api都在OpKernelConstruction這個類裡面建議詳細閱讀。

具體實現過程不細說了，可以直接上C，也可以利用強大的Eigen，裡面有很好的並行化機制根據各種cost來分配資源（但我不會用。。模板元編程那套不學門函數式語言真是入不了門，所以我直接上std::thread了）。GPU的實現我不在行，一般也就拿CUDA_1D_KERNEL_LOOP搞一下，tensorflow裡面利用了大量cub的api很是優雅。

創建python介面

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

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

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

實現op梯度計算

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

測試：同樣官方有示例，推薦一個tf.test裡面的compute_gradient和compute_gradient_error，很好用。

最後祝大家煉丹順利~~找到理想工作。

diss eager

只要是tf寫好的op都能自動求導，那麼只要正向寫好一個層，就能鏈式法則自動求導，跟靜態圖和動態圖無關。

但是話說回來，效率很可能不如手動優化的高

舉例子用的是加法乘法。。那麼你真正想要的一定是定義自己的op。

在tf的首頁http://tensorflow.google.cn，最上方的一排大字里有個extend，裡面有個adding a new op。挺詳細的，照著試一試就好了。

再不濟，百度搜索也能一堆自定義op的文章。雖然大部分是翻譯官方教程，但也有一些是有自己的例子的。

自定義特殊的操作是怎麼樣的呢？可不可以理解為類似C++標準已經提供了一些內建數據類型，然後再自己定義新的ADT？那麼這就相當於組合使用了官方可以自動求導的介面，不就可以達到自動求導了嗎？