選自medium
作者：Josh Fromm
機器之心編譯
機器之心編輯部
去年，TVM 開發團隊陳天奇等人創建了 OctoML 公司，旨在「讓機器學習可以部署在所有硬體上」。前段時間，該公司推出了第一個軟體即服務產品 Octimizer，可以幫助開發者更方便、快捷地將 ML 模型部署到設備上。近日，該公司官方博客又介紹了一種快速端到端二值神經網絡——Riptide，使用 TVM 進行優化時可以實現最高 12 倍的端到端加速。該公司機器學習系統工程師 Josh Fromm 在博客中介紹了 Riptide 的細節。
Riptide 是一種新的模型量化方法，可以將模型量化至 1、2 位。研究團隊今年三月在 MLSys 上介紹了 Riptide，這篇文章主要講一下為什麼要構建 Riptide，並快速了解它的幕後工作原理。團隊計劃來年將 Automatic ultra low-bit 功能添加到 Octomizer 中。在此之前，讀者可以使用開源 Riptide 項目和 MLSys 論文中的信息來進行模型優化。
論文連結：https://proceedings.mlsys.org/static/paper_files/mlsys/2020/155-Paper.pdf
GitHub 項目：https://github.com/jwfromm/Riptide
動機及背景
機器學習發展迅速。幾乎每個月，那些優秀的新模型都會大大提高一些視覺或語言任務方面的 SOTA。這些改進一部分是由新的算法和架構創新所推動的，但對於深度學習任務來說，不斷擴展的算力和內存也使其取得了重大進展。
隨著 ML 準確率的提升，模型所需要的算力和內存也不斷增加。
早在 2016 年，我們就可以看出模型大小和準確率之間的關係了。當下，許多 SOTA 模型只能在最新的 NVIDIA GPU（或 GPU 集群）上才能有效運行。而許多用戶無法花費數千美元在高端 GPU 上，所以就限制了模型部署的範圍。
由於各類網絡可能只需要訓練一次，開發人員也許可以證明這種訓練費用是合理的、有效的。但是，一旦模型進行了部署，將長時間被大量用戶不斷運行。要以這種規模部署最新的模型，通常需要將輸入數據流傳輸到雲上，並將預測返回到用戶設備。現在許多應用程式都依賴這種方法，但它有一些缺點，體現在網絡連接性、延遲、隱私問題以及龐大複雜的基礎架構上。
為了避免這些缺點，許多團隊探索如何在低成本的終端用戶硬體（如手機或 IoT 設備）上直接運行最新模型。這是一個巨大的挑戰，因為此類設備沒有足夠的算力或內存。例如，樹莓派 3 比 NVIDIA Titan GPU 慢了大約 4000 倍。
為了使高精度模型適應此類平臺，最近的研究方向已經開始探索如何使這類網絡運行更快，同時佔用更少的內存。從較高的層面來說，這些技術遵循兩種策略：體系架構優化和近似優化。架構優化涉及尋找連接層的新方法，以減少延遲或提高參數有效性。MobileNet 和 SqueezeNet 是兩個以移動端為重點的體系架構。與創建新的行動裝置友好型模型相反，近似優化旨通過加快運算速度提高現有模型的速度，同時保持足夠的準確率。
兩種流行的近似優化方法。
在最近流行的幾種近似優化中，知識蒸餾和剪枝是兩種有代表性的方法（如上圖所示）。前者嘗試使用大型教師網絡來更好地訓練學生網絡；後者則刪除了網絡中影響較小的權重和激活函數。本文將重點介紹另外一種方法——二值神經網絡，這種網絡在泛化性能、加速潛力、內存壓縮等方面都有優秀的表現。
二值網絡
為了提高性能並減少內存需求，研究團隊在部署模型時越來越多地量化激活函數和權重。例如，在推理過程中，工程師可能將模型轉換為能夠對 int8（8 位整數）進行運算的方式，而不是在訓練過程中經常使用的 float32（IEEE 754 單精度浮點數）。小型整數運算不僅比浮點運算速度快，更因為它們佔用的位數更少，所以可以通過充分利用可用內存帶寬來提高吞吐量。在實踐中，很多模型被量化後都不會有顯著的準確度損失，因此這項技術大受歡迎。
二值化（binarization）將量化進行到了極致，將網絡的權重和激活函數降低到只剩一位，這會帶來一些新的優化。考慮兩個 1 位值之間所有可能的乘法（如下圖），這類方程式和「與門」的邏輯真值表極為相似。
如果認為值為 0 的位代表-1，則上表將成為「同或門」的真值表。
這種等效可以讓我們用更高效的二值運算代替浮點運算。
浮點點積（上）和二值點積（下）的比較。
對上圖點積內循環中的運算次數進行統計可以發現，二值化可以將一層中的運算次數減少到原來的 1/43，將參數大小減少至 1/32。當這樣大體量的優化效果首次由 Courbarioux 等人提出時，2016 年便掀起了二值網絡研究的熱潮。當然，用 1 位值近似 32 位浮點數是一種有損近似。與對應的全精度相比，二值網絡通常會有明顯的精度損失，top-1 準確率通常會損失近 20％。因此，二值網絡的研究重點一直聚焦於如何減少精度損失。
儘管多個研究團隊在提高二值網絡準確率方面取得了巨大進步，但他們都沒有以一種可以衡量端到端加速的方式實現該網絡。這種實現的缺乏不僅讓我們很難知道二值網絡的實際速度，而且還阻止了二值網絡在許多實際環境中的應用。
而缺乏實現的原因是，要編寫一個簡單的二值矩陣乘法，儘管運算次數比較少，但它依然會比大多數的全精度乘法慢得多。函數的運行時間並不完全取決於其運行次數，內存的訪問方式也起著重要作用。任何二值網絡的實現都必須與 OpenBLAS 和 MKLDNN 這樣的庫競爭，而這些庫經歷了大型工程團隊多年的手工優化，難度可行而知。而對於大多數研究機構來說，花大量的時間和精力來建立具有競爭力的庫根本不可能。取而代之的是，他們在訓練期間模擬二值化，並假設根據運算次數可以預測加速。
Riptide 的突圍
為了解決這些問題，OctoML 的研究者提出了 Riptide，這是一種找出並解決端到端二值網絡瓶頸的方法。Riptide 基於 TVM，後者是一種深度學習系統編譯器，可以幫助我們自動生成經過調優的高性能二值化算子。
迄今為止，二值網絡優化僅僅著眼於高效實現核心低位（low-bit）卷積層的不同策略。這一著眼點基於以下假設，即二值網絡性能可以反映高精度網絡的行為：如果核心卷積可以利用儘可能少的位數實現足夠高的準確率，那麼整個網絡將變得非常快。然而，沒有哪個二值網絡是單純地由卷積構成。在卷積之間有許多關鍵的中間運算，需要用它們來處理下一層的數據。在高精度網絡中，這些層的延遲可以忽略不計，但在二值網絡中，卷積可以實現 43 倍的加速，這些中間「粘合層（glue layers）」就變得非常重要。
二值卷積之間的「粘合層」及其計算複雜度。H 和 W 表示輸入維度，F 表示濾波器的數量。
當前多數的二值網絡至少包含 4 個上圖中的藍色層（QConv 和 Quantize 之間）。從左到右來看，首先將 QConv 的輸出從整數形式去量化為等效的浮點數，接下來，使用權重縮放（傳播實值權重的大小）和批歸一化（保持激活的分布可預測）進行縮放。然後應用 ReLU 等非線性激活函數，結果被重新量化為單個位，並打包為下一個量化卷積做準備。
為了理解粘合層的重要性，可以聯想一下 SqueezeNet（一種更高效的移動端部署架構）。假設典型的輸入大小約為 200x200 像素，並且二值卷積的完美 roofline 實現運行速度是全精度的 43 倍，那麼可以估計出網絡在粘合層中花費的總執行時間。
假設二值化可以使卷積的速度提高近 43 倍，可以進一步估計，更高精度的粘合層將消耗大約 70% 的總推理時間。這是一個相當大的瓶頸！即使在較低的卷積速度（如 20 倍和 10 倍），粘合層仍然會消耗大約一半的推理時間。因此，研究者認為，要真正實現二值網絡所承諾的加速，粘合層也必須是二值化的!
前面已經介紹過融合粘合運算（fused glue operation），它只用兩個指令就完全替代了粘合層。細節可以參見論文。其關鍵思想是用移位運算代替乘法，將縮放項近似為 2 的近似冪，用定點量化近似（fixed point quantized approximations）代替浮點加法和減法。定點量化近似可以直接添加到二值卷積輸出中。綜上所述，得到如下方程組：
其中，N 是用來量化網絡激活的位數，最後一行求解 q(a) 給出了融合粘合的完整方程。通過替換這種融合粘合運算，可以創建一個完全二值化的網絡：
研究者還在 ImageNet 數據集上進行了廣泛的準確度掃描，發現與其他 SOTA 二值網絡相比，上文中的融合粘合運算不會造成任何準確度損失。
這是一個好消息，因為在體系架構層次上消除了粘合瓶頸之後，現在已經可以著手生成快速的二值卷積並度量端到端性能。
端到端加速
早期，研究者將二值網絡作為一種近似技術，幫助我們在移動和 IoT 設備上部署高效的模型。先來看一下 Riptide 在樹莓派 3b 上的表現如何。樹莓派 3b 基於 ARM Cortex-A53 處理器，是資源受限環境的二值目標的典型代表。
首先，將基於 for 循環的簡單網絡實現與使用 MKLDNN 來加速其浮點運算的完整精度基線進行比較。
使用 MKLDNN 的完整精度 ResNet18 運行時間 vs. 未優化的完全二值網絡。
在沒有優化的情況下，研究者獲得了 ResNet18 的一個完全二值網絡（FBN）版本，其速度大約相當於對應高精度網絡的 1/5。為了真正利用二值網絡，現在還需要生成一個經過調優的執行 schedule。研究者通過在 TVM 中實現二值算子並利用其強大的調度、優化能力來得到這種 schedule。
TVM 提供了以下可以用於優化函數的 schedule 特性:
Tiling 將一個計算分解成多個塊，以改善負載的內存局部性。
Vectorization 利用硬體 SIMD 指令來實現更高效的運算執行。
Parallelization 利用多核等 MIMD 設施。
Loop Unrolling 複製循環體來減少開銷。
此外，研究者還利用了《Automating Generation of Low Precision Deep Learning Operators》論文中提到的 fast popcount 算子。最後，他們還介紹了一種名為 bitpack fusion 的新優化方法。
Bitpack fusion 將 bitpacking 摺疊成卷積，以減少激活內存消耗。
在高級層次上，bitpack fusion 儘可能將 bitpacking 摺疊到前面的卷積核中。這可以將中間內存需求減少到原來的 1/16，使得複製策略更加高效，從而提供進一步的加速。
那麼，這一堆優化究竟哪些更重要呢？
為了解答這一問題，研究者在 SqueezeNet 上進行了控制變量研究，結果如上圖所示。
他們發現，與更高精度的浮點基線相比，每種優化都對加速有顯著影響。如果將各種優化方法一起應用，可以在真實模型上看到 10 倍的加速。雖然 10 倍的加速遠小於論文中所期望的 43 倍加速，但這也是一個數量級的提升，可以顯著擴展能夠高效運行 SOTA 模型的設備範圍。當然，這種加速只有在保證準確率的前提下才有價值。
上表顯示了幾種流行二值模型在 1、2、3 位激活條件下的準確率與運行時間權衡。總體來看，可以發現在提供更多加速的同時，Riptide FBN 和其他 SOTA 二值技術一樣穩定，甚至比後者更加準確。Riptide FBN 可以提供準確的激活位寬和量化極性（quantization polarity），提供 4~12 倍的加速，使得它很容易滿足所有應用的準確性需求。
原文連結：https://medium.com/octoml/riptide-fast-full-binarization-in-tvm-ae2afd2104bb