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

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

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

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

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

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

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

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

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

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

在此,我們將借用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暫時還沒有包括時序模型。

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

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

1。2 池化層

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

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

1。3 批資料歸一化層

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

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

1。4。 資料歸一化

將資料進行歸一化處理,通常用於輸出層的歸一化。

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

1。5。 其他計算層

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

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

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暫時不支援。兩者都對指數運算和對數運算,開方運算提供了支援,可是這兩個函式其實是加速晶片比較難精確實現的函式(可以透過查詢表或近似函式實現)。

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

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目前還缺乏對通用矩陣運算的支援。

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

2。3 啟用和非線性函式

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

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

2。4。 隨機數和常數

這些操作用於產生資料,包括正態隨機產生,均勻隨機產生,常數等。可以看出ONNX支援隨機數加入到圖中,而NNVM目前還不支援圖中包括隨機數。

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

2。5。 降維繫列

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

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

3. Tensor變換

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

3。1。 分割組合運算元

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

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

3。2。 索引變換

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

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

3。3。 資料選取

此部分操作可以根據維度引數、邊框或者腳標矩陣引數選取Tensor的部分資料,或者對Tensor的資料進行復制拓展。

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

3。4。 資料填充

資料填充分為邊緣補0,常數填充和複製。其中NNVM沒有在官方文件頁面中提供fill函式的解釋,但是確實存在這個函式。

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

綜上,我們總結了作為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/bl

ob/master/docs/Operators。md

[3] NNVM Core Tensor Operators,

http://

nnvm。tvmlang。org/top。ht

ml#detailed-definitions

[4] NETRON,

https://

github。com/lutzroeder/N

etron

作者簡介:

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

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

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