深度學習飛速發展過程中,人們發現原有的處理器無法滿足神經網絡這種特定的大量計算,大量的開始針對這一應用進行專用芯片的設計。谷歌的張量處理單元(Tensor Processing Unit,后文簡稱TPU)是完成較早,具有代表性的一類設計,TPU采用基於脈動陣列設計的矩陣計算加速單元,可以很好的加速神經網絡的計算。本系列文章將利用公開的TPU V1相關資料,對其進行一定的簡化、推測和修改,來實際編寫一個簡單版本的谷歌TPU,以更確切的了解TPU的優勢和局限性。
動手寫一個簡單版的谷歌TPU系列目錄
拓展
TPU的邊界(規划中)
重新審視深度神經網絡中的並行(規划中)
1. 完成SimpleTPU的設計
在 谷歌TPU概述和簡化中給出過SimpleTPU的框圖,如下圖所示。
在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的各個運行單元可以流水並行的,即
而針對網絡中計算量最大的全連接層和卷積層,針對性設計的乘法整列和向量計算的設計方法可以讓其在每個時鍾周期都完成有效的乘加計算;這意味着和CPU相比,SimpleTPU可以達到極高的效率。