Simple TPU的設計和效能評估

暗海風發表於2019-07-26

深度學習飛速發展過程中,人們發現原有的處理器無法滿足神經網路這種特定的大量計算,大量的開始針對這一應用進行專用晶片的設計。谷歌的張量處理單元(Tensor Processing Unit,後文簡稱TPU)是完成較早,具有代表性的一類設計,TPU採用基於脈動陣列設計的矩陣計算加速單元,可以很好的加速神經網路的計算。本系列文章將利用公開的TPU V1相關資料,對其進行一定的簡化、推測和修改,來實際編寫一個簡單版本的谷歌TPU,以更確切的瞭解TPU的優勢和侷限性。

動手寫一個簡單版的谷歌TPU系列目錄

 

    谷歌TPU概述和簡化

    TPU中的脈動陣列及其實現

    神經網路中的歸一化和池化的硬體實現

    TPU中的指令並行和資料並行

    Simple TPU的設計和效能評估

    SimpleTPU例項:影象分類

    擴充

    TPU的邊界(規劃中)

    重新審視深度神經網路中的並行(規劃中)

 

1. 完成SimpleTPU的設計

    在 谷歌TPU概述和簡化中給出過SimpleTPU的框圖,如下圖所示。

clip_image002

    在TPU中的脈動陣列及其實現中介紹了矩陣/卷積計算中的主要計算單元——乘加陣列(上圖4),完成了該部分的硬體程式碼並進行了簡單的驗證;在 神經網路中的歸一化和池化的硬體實現中介紹了卷積神經網路中的歸一化和池化的實現方式(上圖6),同時論述了浮點網路定點化的過程,並給出了Simple TPU中重量化的實現方式,完成了該部分的硬體程式碼並進行了驗證。

    在 TPU中的指令並行和資料並行中對整個處理單元的體系結構進行了分析和論述,包括指令並行和資料並行兩個方面。那麼,如何在TPU中的指令並行和資料並行中提到的設計思路下,將TPU中的脈動陣列及其實現神經網路中的歸一化和池化的硬體實現中提到的計算單元充分的利用,是完成Simple TPU設計的最後一部。根據SimpleTPU的框圖可知,需要實現的功能包括

  • 指令的取指和譯碼(上圖4)
  • Weight的讀取(上圖2)
  • 各個執行單元的控制和排程(上圖1)
  • 讀取影象和寫回結果(上圖5)

    在SimpleTPU的設計中,指令的取指和譯碼和Weight的讀取功能都較為簡單,可直接參照程式碼。

    在對各個執行單元進行控制和排程時需要確保各個單元可以共同執行,沒有相互之間的資料以來關係。

    除此之外,還需要單獨實現讀取影象和寫回結果的功能。SimpleTPU中只關注核心的計算功能,該部分功能並未進行優化,後續對實現效果進行分析時,也會將該部分執行時間排除在外。

    至此,Simple TPU的設計基本完成了,程式碼可參見https://github.com/cea-wind/SimpleTPU

2. SimpleTPU的特性

    SimpleTPU的主要特性包括

  •     支援INT8乘法,支援INT32的累加操作
  •     採用VLIW進行指令並行
  •     採用向量體系結構進行資料並行

    SimpleTPU依照Google TPU V1的設計思路,可以完成神經網路推理過程中的大部分運算。依據設計,支援的運算包括(理論)

運算

說明

Conv3d

in_channels:資源受限

out_channels:資源受限

kerner_size:幾乎無限制

stride:幾乎無限制

padding:幾乎無限制

dilation:幾乎無限制

groups:極有限支援,架構限制

bias:支援

ConvTranspose3d

同上

Maxpool2d

kernel_size:幾乎無限制

stride:幾乎無限制

padding:幾乎無限制

Avgpool2d

同上

Relu

僅支援RELU作為非線性函式

BatchNorm2d

推理中BatchNorm2d被融合到Conv或者Pool完成

Linear

資源受限

UpscalingNearest2D

多次呼叫pool完成

UpscalingBilinear2D

多次呼叫avgpool完成

     

    其中,資源受限代表該引數的取值範圍有限,主要受限於SimpleTPU的儲存設計等;幾乎無限制表示其取值僅受到暫存器位寬等限制。由於架構設計上的問題,SimpleTPU對groupconv支援極為有限,在不合適的引數下效率可能遠低於普通卷積;類似的,Google TPU也不能很好支援groupconv,並明確告知不制止depthwise conv(極度稀疏化的group conv)。

    BatchNorm2d在推理過程中實際上時進行逐點的乘法和加法,其中加法計算可以融合到下一層或者上一層的卷積計算中進行,乘法計算可以和pooling計算融合。在SimpleTPU設計,Pooling模組實際上一直在工作,即使網路中沒有pooling層,SimipleTPU增加了一個1*1,stride=1的pooling層進行等價。

    Upscaling操作通過pooling完成計算。這是因為在SimpleTPU中,reshape操作(支援的)是沒有代價的。pooling操作可以完成雙線性插值的計算,因此可以完成upscaling中的所有數值的計算。可以理解為通過pooling+reshape完成了upscaling的計算。

   

3. SimpleTPU的效能

    Simple TPU設計了一個32×32的int8乘加陣列計算矩陣乘法和卷積,和一個1×32的int32乘法陣列進行池化和歸一化的計算。根據Xilinx HLS工具的綜合結果,在UltraScale+系列的FPGA器件上,工作頻率可達500MHz。因此SimpleTPU的算力約為

32×32×500MHz×2 = 1Tops

    作為對比,GoogleTPU V1的算力約為91Tops(int8),差異主要在SimpleTPU的規模為其1/64,同時在FPGA上的工作頻率會低於ASIC的工作頻率。

    依據設計,SimpleTPU在適合的任務下會有很高的執行效率,TPU中的指令並行和資料並行中針對這一點又更為具體的描述。從巨集觀上看,SimpleTPU的各個執行單元可以流水並行的,即

image

    而針對網路中計算量最大的全連線層和卷積層,針對性設計的乘法整列和向量計算的設計方法可以讓其在每個時鐘週期都完成有效的乘加計算;這意味著和CPU相比,SimpleTPU可以達到極高的效率。

相關文章