內存帶寬與計算能力,誰才是決定深度學習執行性能的關鍵?
說到模型對于硬件的要求,大家***個想到的就是計算量,即一個深度學習模型需要多少次計算才能完成一次前饋。然而,除了運算量之外,模型對于內存帶寬的需求也是影響實際計算所需要時間的重要參數。我們下面會看到,在內存帶寬有限的情況下,僅僅縮小計算量并不能讓計算時間等比例下降!
內存帶寬對于硬件系統的性能影響如上圖所示。如果把內存比做瓶子,運算單元比作杯子,那么數據就是瓶子里的各色顆粒,而內存接口就是瓶口,通過瓶口數據才能進入杯子被消費(處理)掉。而內存帶寬就是瓶口的寬度了。瓶口寬度越窄,則數據需要越多時間才能進入杯子(處理單元)。正所謂「巧婦難為無米之炊」,如果帶寬有限,那么即使處理單元***快,在大多數時候也是處理單元在空等數據,造成了計算力的浪費。
深度學習網絡與 Roofline 模型
對于工程師來說,定性分析并不夠,我們還需要能定量分析算法對于內存帶寬的需求,以及對于計算性能的影響。
算法對于內存帶寬的需求通常使用「運算強度 (operational intensity,或稱 arithmetic intensity)」這個量來表示,單位是 OPs/byte。這個量的意思是,在算法中平均每讀入單位數據,能支持多少次運算操作。運算強度越大,則表示單位數據能支持更多次運算,也就是說算法對于內存帶寬的要求越低。所以,運算強度大是好事!
我們來舉一個例子。對于步長(stride)為 1 的 3x3 卷積運算,假設輸入數據平面大小為 64x64。簡單起見,假設輸入和輸出 feature 都為 1。這時候,總共需要進行 62x62 次卷積運算,每次卷積需要做 3x3=9 次乘加運算,所以總共的計算次數為 34596,而數據量為(假設數據和卷積核都用單精度浮點數 2byte):64x64x2(輸入數據)+ 3x3x2(卷積核數據)= 8210 byte,所以運算強度為 34596/8210=4.21。如果我們換成 1x1 卷積,那么總的計算次數變成了 64x64=4096,而所需的數據量為 64x64x2 + 1x1x2=8194。顯然,切換為 1x1 卷積可以把計算量降低接近 9 倍,但是運算強度也降低為 0.5,即對于內存帶寬的需求也上升了接近 9 倍。因此,如果內存帶寬無法滿足 1x1 卷積計算,那么切換成 1x1 卷積計算雖然降低了接近 9 倍計算量,但是無法把計算速度提升 9 倍。
這里,我們可以看到,深度學習計算設備存在兩個瓶頸,一個是處理器計算能力,另一個是計算帶寬。如何分析究竟是哪一個限制了計算性能呢?可以使用 Roofline 模型。
典型的 Roofline 曲線模型如上圖所示,坐標軸分別是計算性能(縱軸)和算法的運算強度(橫軸)。Roofline 曲線分成了兩部分:左邊的上升區,以及右邊的飽和區。當算法的運算強度較小時,曲線處于上升區,即計算性能實際被內存帶寬所限制,有很多計算處理單元是閑置的。隨著算法運算強度上升,即在相同數量的數據下算法可以完成更多運算,于是閑置的運算單元越來越少,這時候計算性能就會上升。然后,隨著運算強度越來越高,閑置的計算單元越來越少,***所有計算單元都被用上了,Roofline 曲線就進入了飽和區,此時運算強度再變大也沒有更多的計算單元可用了,于是計算性能不再上升,或者說計算性能遇到了由計算能力(而非內存帶寬)決定的「屋頂」(roof)。拿之前 3x3 和 1x1 卷積的例子來說,3x3 卷積可能在 roofline 曲線右邊的飽和區,而 1x1 卷積由于運算強度下降,有可能到了 roofline 左邊的上升區,這樣 1x1 卷積在計算時的計算性能就會下降無法到達峰值性能。雖然 1x1 卷積的計算量下降了接近 9 倍,但是由于計算性能下降,因此實際的計算時間并不是 3x3 卷積的九分之一。
顯然,一個計算系統的內存帶寬如果很寬,則算法不需要運算強度很大也能輕易碰到計算能力上限決定的「屋頂」。在下圖中,計算能力不變,而隨著內存帶寬的上升,達到計算力屋頂所需的運算強度也越低。
Roofline 模型在算法-硬件協同設計中非常有用,可以確定算法和硬件優化的方向:到底應該增加內存帶寬/減小內存帶寬需求,還是提升計算能力/降低計算量?如果算法在 roofline 曲線的上升區,那么我們應該增加內存帶寬/減小內存帶寬需求,提升計算能力/降低計算量對于這類情況并沒有幫助。反之亦然。
我們來看一個實際的例子,比較一下各種機器學習算法在 roofline 模型上所處的位置。下圖取自 Google 的 TPU 論文《In-Datacenter Performance Analysis of a Tensor Processing Unit》。由圖中可見,LSTM 算法的運算強度***,所以被卡在了 roofline 模型的上升區中間的地方,即 TPU 在執行 LSTM 算法的時候,由于內存帶寬限制所以性能只有 3TOPS 左右,僅為峰值性能(90TOPS)的三十分之一。經典全聯接神經網絡(multi-layer perceptrons, MLP)的運算強度略好于 LSTM,也被卡在 roofline 曲線的上升區,實際執行性能大約在 10TOPS 左右。而卷積神經網絡模型,尤其是 CNN0,由于卷積神經網絡中能實現卷積核復用,因此運算強度非常高,于是可以非常接近 TPU roofline 曲線的屋頂(86 TOPS)。CNN1 模型雖然運算強度也很高,但是由于種種其他原因(論文中表示是由于 CNN1 模型的特征深度較淺無法完全利用 TPU 的計算單元)無法到達屋頂。這個例子又讓我們看到了硬件-算法協同設計時的另一個要點:除了內存帶寬之外還有「其他原因」可能讓算法無法到達屋頂,我們要盡量減小這些「其他因素」!