從NNVM和ONNX看AI晶元的基礎運算運算元

AI晶元的需求分析是以應用和演算法為出發點的，深度學習框架作為演算法設計的入口，是非常好的參考。但目前的深度學習框架種類繁多，提供的功能也比較複雜，不利於直接分析硬體需求。相比之下，中間表示層（IR）往往更接近硬體實現，可以認為是軟硬體之間的橋樑，對AI晶元的專用指令集設計具有非常高的參考價值。因此，我們不妨從兩個中間表示層工具（NNVM和ONNX）出發，通過觀察它們的原語操作，思考我們需要在AI晶元中支持什麼樣的運算運算元。（關於中間表示層IR的一些基本概念可以參考Deep Learning的IR「之爭」）

NNVM是由陳天奇團隊提出的一套可復用的計算流圖中間表達層，它提供了一套精簡的API函數，用以構建、表達和傳輸計算流圖，從而便於高層級優化。另外NNVM也可以作為多個深度學習框架的共享編譯器，可以優化、編譯和部署在多種不同的硬體後端。其特點是部署的模型擁有最小依賴，可加入新的操作(operators)，可將新的優化通路加入到現有的圖結構中。從NNVM的觀點看，它可用於將Ｍ個框架，Ｎ個機器之間構建一個單一的紐帶，以實現各種框架向各種實現平台的無差別部署。如圖所示。

ONNX[1]是Facebook聯合微軟和AWS推出的開源的深度學習表示格式。通過ONNX，AI開發人員可以容易地在不同模型和工具間轉換，並將工具組合使用。目前可以支持Caffe2, CNTK, PyTorch等主流模型，也對Tensorflow提供了早期支持。一下是ONNX的一些基本信息（來自ONNX網站[1]）。

嚴格而言，和NNVM相比，ONNX更像是一個協議轉換器，可以在各個框架之間進行轉換。

目前多種標準的轉換工具已經開發了出來，如表格所示。

在此，我們將借用NNVM和ONNX的運算元，分析AI硬體加速的需求。這些運算元都包含了相關參數，在此不細緻表述。有關這些運算元的完整解釋，可以查看ONNX的官方表述[2]；NNVM目前的唯一版本是0.8，因此本文參考版本0.8。運算元的官方解釋參見[3]。

ONNX以張量(Tensor)作為輸入輸出數據格式，可以認為是多維數組。主要支持如下的數據格式：tensor(float16), tensor(float), tensor(double)。目前還不太支持定點，但提供了定點化的一些函數，如floor, ceil, clip, cast等。NNVM也以Tensor作為數據格式。

ONNX不完整支持神經網路的訓練，但它提供了訓練所需要的完整的圖描述，其實就是對BN和dropout兩個功能上加了一個區別參數，用來描述可訓練的圖和不可訓練的圖。由於ONNX實際上就是把各種框架的圖轉換成了它自己的Operator表示的圖，它只負責描述這個圖的結構，具體前後向的計算都需要一個執行框架（稱為後端）。 因此，如果需要實現訓練，需要實現系統根據這個圖自動實現所有的反向過程的微分運算元，並且實現loss函數。此外，還需要編譯器根據這個前向圖結構推演出反向圖的各個步驟，這個過程可以是工具鏈自動的。而對於前向而言，基本不需要轉化就可以直接部署在支持這些運算元的實現平台上。

本人對這些運算元進行了一個簡單的歸類，並將ONNX和NVM進行了比對。由於在硬體實現上，不同的運算元的實現複雜度是不同的，因此加入了Complexity的度量。另外，根據當前神經網路在圖像/語音/文本三方面的應用情況，對這些運算元的使用頻率進行了估計。由於應用領域和硬體平台各不相同，因此複雜度和使用頻率僅作參考。

1. 深度神經網路計算

1.1. 計算層

這部分運算元是深度網路的核心，用於將輸入的神經元激活值與突觸連接強度（權重）進行積分求和，得到新的神經元的模電位。根據是否滑窗，是否具有時序結構，可分為如下幾種運算元，其中FC（全連接）是多層感知機(MLP)的基礎，Conv和FC是深度卷積神經網路的基礎。RNN, GRU, LSTM是帶有時序結構的神經網路模型，主要用於非靜態圖像的場合，例如語音/文字/視頻等。可見，ONNX的關注面比較全面，包括了時序模型，而NNVM暫時還沒有包括時序模型。

注: (*)代表ONNX庫中此函數帶有實驗階段(Experimental)標誌。下同。

1.2 池化層

池化層主要用於尺度變換，提取高維特徵。主要分三種池化，第一種是局部池化，在圖像維度上，幾個相鄰的點被縮減為一個輸出點，在Channel維度上不變。包括計算相鄰點的平均值（AveragePool），最大值（MaxPool），範數（LpPool）。主要用於圖像的尺寸變換。第二種是全局的池化，此時一個Channel的所有數據點縮為1個點，因此有幾個Channel就輸出幾個數據點。此種類型主要用於深度卷積神經網路中卷積部分與FC部分的連接。第三種是ROI-pooling，用於Faster-RCNN等檢測識別網路中對感興趣區域進行尺度歸一化，從而輸入到識別網路進行識別處理。可見，ONNX實現了比較全面的運算元覆蓋，NNVM實現了比較常見的局部池化和全局池化，但是暫時還沒有實現ROI-pool。

1.3 批數據歸一化層

歸一化層作為一個特殊層，可用於數據的歸一化，提高神經網路的性能，降低訓練時間。對於帶有殘差的神經網路非常重要。目前高性能網路大多帶有歸一化層，而絕大多數都會採用Batch Normalization(BN)。BN前向操作並不複雜，但反向比較複雜，因此用於訓練的BN需要加入更多的子層。ONNX構建了兩套圖描述，用標誌位進行區分，用戶可以選擇是用於訓練的還是僅用於前向的。另外，ONNX還提供了在這方面其他的選擇，例如Instance歸一化 (y = scale * (x - mean) / sqrt(variance + epsilon) + B)和基於範數的歸一化，LRN被用在AlexNet等早期設計，目前用的比較少。對比之下，NNVM只支持了BN，可以覆蓋約95%的應用情形。

1.4. 數據歸一化

將數據進行歸一化處理，通常用於輸出層的歸一化。

1.5. 其他計算層

在進行訓練時，DropOut隨機扔掉一些通路，可以用於防止過擬合。這方面兩個框架都實現了。Embedding用於將詞轉換為高維表達，是文本的預處理的主要步驟。GRUUnit是個試驗性函數，功能類似於GRU的激活層。

2. 基礎Tensor運算

2.1 逐元素運算（element-wise）類

這個類別包括了Tensor的一些基礎運算，由於輸出的數據點只跟對應的那一個輸入的數據點有關係，因此可以稱為element-wise運算，這類運算與輸入的數據的維度和結構無關，可以等價的認為是一維向量運算的Tensor等效表示。由於輸入數據可能是各種維度，也可以是標量，因此此中的操作都是維度兼容的。一種特殊情況是輸入的參數中，一個是向量，另一個是標量，此時，NNVM區別對待，而ONNX將其統一處理。在此處，ONNX只支持了Tensor-scalar，不支持scalar-tensor, 而NNVM兩者都支持。除法情況類似。對於加減乘除，ONNX自帶broadcast操作，而NNVM通過單獨函數實現。ONNX對邏輯運算和比較運算提供了支持，而NVVM沒有。另外ONNX提供了一些數據格式轉換(cast)和量化方面的函數(clip, floor, ceil)，而NVVM暫時不支持。兩者都對指數運算和對數運算，開方運算提供了支持，可是這兩個函數其實是加速晶元比較難精確實現的函數(可以通過查找表或近似函數實現)。

2.2 Tensor/矩陣處理類

這部分操作是對整個Tensor的數據進行的，即輸出可能關係到Tensor中的不止一個數據。包括求和，求平均，通用矩陣運算(Gemm)，矩陣乘法，圖像縮放等。其中Gemm是矩陣處理的通用表達形式，即Y = alpha * A * B + beta * C。其中A 為M X K維, B為K X N維，C和Y為M X N維。可以認為NNVM目前還缺乏對通用矩陣運算的支持。

2.3 激活和非線性函數

激活函數提供了神經網路的非線性擬合能力，不同的激活函數具有各自的性能特點。由於ReLU簡單且性能較好，因此一般圖像處理演算法採用ReLU函數。而Sigmoid和Tanh在LSTM/GRU/RNN中較為常見。這些函數可以認為是2.1中所述的element-wise運算元，但為了表達其在神經網路中的特殊功用，在此單獨提出。

2.4. 隨機數和常數

2.5. 降維繫列

降維繫列是ONNX特有的。可以指定哪些維度根據某個計算度量去除。由於計算度量的方法比較多，本人認為功能類似於Global pooling。這些方法在神經網路應用中不是很多，NNVM目前還不支持這些方法。

3. Tensor變換

此部分運算元不會改變Tensor的數據，只會對數據的位置和維度進行調整。

3.1. 分割組合運算元

此部分可以將多個Tensor合併成一個，或者將一個拆分為多個。可以用於分組卷積等。

3.2. 索引變換

索引變換包括Reshape, 矩陣轉置，空間維度與Feature Map互換等。可以認為是數據排布關係的變化。Flatten和Squeeze可以認為是Reshape的特例。

3.3. 數據選取

此部分操作可以根據維度參數、邊框或者腳標矩陣參數選取Tensor的部分數據，或者對Tensor的數據進行複製拓展。

3.4. 數據填充

數據填充分為邊緣補0，常數填充和拷貝。其中NNVM沒有在官方文檔頁面中提供fill函數的解釋，但是確實存在這個函數。

綜上，我們總結了作為IR表示層的所有操作(Operator)。將這些操作連接起來就構成了數據流圖，使得神經網路可以表達為一個基於Operator和Tensor的有向圖。採用Netron[4]可以查看ONNX的數據流圖，具有很好的可視化體驗。推薦大家可以嘗試。不過NNVM目前好像還沒有類似的工具。

另外，總體感覺是ONNX的靈活度高於NNVM，尤其是在RNN的支持上邊，但NNVM給大家提供了一個很好的範例，用以說明如何抓住重點，覆蓋典型的應用場景。另外，NNVM提供了一個很好的擴展機制，用戶可以將自己的原子操作加入到框架中去而不改變原有的框架結構。上文中提及的可以認為是本徵支持的操作（native operation）。

對於一個視頻/圖像類神經網路晶元，可以先考慮支持NNVM所支持的本徵原語部分，如果有需要，再向著ONNX的更多操作擴展。而對於文本/語音處理而言，ONNX是比較好的考評量度。也許將來大家在進行加速器的功能比拼時，會以ONNX框架提供的兼容性為尺度。

另外，值得注意的是，要做到對一個操作的支持，並不只是有沒有的問題，還包括執行效率的問題。後者可以從功耗效率的角度來衡量，也可以從有效計算能力和峰值計算能力的比值看出來。因此想設計一個高靈活的神經網路晶元還是一個各方面權衡，靈活度與性能聯合調優的過程。

[ Reference ]

[1] Open Neural Network Exchange (ONNX), https://onnx.ai/

[2] ONNX Operator Schemas, https://github.com/onnx/onnx/blob/master/docs/Operators.md

[3] NNVM Core Tensor Operators, http://nnvm.tvmlang.org/top.html#detailed-definitions

[4] NETRON, https://github.com/lutzroeder/Netron

作者簡介：

吳臻志博士，清華大學類腦計算研究中心助理研究員。專長神經網路晶元設計，眾核晶元設計，神經網路高效實現等。郵件 wuzhenzhi@gmail.com。

歡迎關注我的微信公眾號:StarryHeavensAbove

題圖來自網路，版權歸原作者所有

推薦閱讀：