內存帶寬與計算能力，誰才是決定深度學習執行性能的關鍵？

機器之心專欄

作者：李飛

隨著深度學習的不斷發展，計算能力得到了深度學習社區越來越多的注意。任何深度學習模型，歸根到底都是需要跑在設備上的，而模型對設備性能的要求越低，則能得到越多的運用——千萬不能讓硬體成為了模型普及的瓶頸！

說到模型對於硬體的要求，大家第一個想到的就是計算量，即一個深度學習模型需要多少次計算才能完成一次前饋。然而，除了運算量之外，模型對於內存帶寬的需求也是影響實際計算所需要時間的重要參數。我們下面會看到，在內存帶寬有限的情況下，僅僅縮小計算量並不能讓計算時間等比例下降！

內存帶寬對於硬體系統的性能影響如上圖所示。如果把內存比做瓶子，運算單元比作杯子，那麼數據就是瓶子里的各色顆粒，而內存介面就是瓶口，通過瓶口數據才能進入杯子被消費（處理）掉。而內存帶寬就是瓶口的寬度了。瓶口寬度越窄，則數據需要越多時間才能進入杯子（處理單元）。正所謂「巧婦難為無米之炊」，如果帶寬有限，那麼即使處理單元無限快，在大多數時候也是處理單元在空等數據，造成了計算力的浪費。

深度學習網路與 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 的計算單元）無法到達屋頂。這個例子又讓我們看到了硬體-演算法協同設計時的另一個要點：除了內存帶寬之外還有「其他原因」可能讓演算法無法到達屋頂，我們要盡量減小這些「其他因素」！

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！