TensorFlow 8 bit模型量化

本文基本參考自這篇文章：8-Bit Quantization and TensorFlow Lite: Speeding up mobile inference with low precision

首先來一段keras dalao Francois Chollet的雞湯：

• make it possible
• make it work
• make it efficient
• make it dependable and invisible
• move on to next layer and think about it again

這段雞湯已經把8-bit量化的需求和階段已經說得很清楚了：提高算法效率

What is 8 bit

目前DL中大部分都是用32bit float類型進行計算的，bit位數的多少直接限制了數據類型可以表達的數據範圍，好比float 32的數據是由1bit表示符號，8bit表示整數部，23位表示分數部組成。

num bits	Min value	Max value
8 bit	-128 (-2^7)	128 (x^7)
16 bit	-32768 (-2^15)	32768(2^15)
32 bit	–2147483648 (-2^31)	–2147483648 (2^31)

用更低位的數值類型意味着更小的數據表示範圍和更稀疏的數值，量化的時候就會形成數值精度損失。好比要把float數值量化到int類型，那麼首先小數部分會損失，而那些超過int類型可以表達的範圍的值也會被壓縮到int可以表達的最大或最小值。

Why 8 bit？

那麼既然會有精度損失那爲啥如今int8運算愈來愈流行，主要緣由就是快跟省了。

• 快：低bit位的數值計算通常會比高bit的要快。雖然現代計算芯片上浮點型的計算已經並不比int類型慢，可是這主要是設計了專用的float計算核，而如今很流行的ai芯片和一些嵌入式芯片上通常並不會設計不少的float計算核，所以對float數值的計算算力很低算力。

  下圖就是Nvidia RTX2080ti的芯片架構(ref)

• 省：從32bit壓縮到8bit，最直接的就是內存可以減小1/4。同時，從RAM中讀取數據的時間也會縮短；也能下降運算能好。

這也就說明了爲何如今愈來愈多的ai芯片專門針對int8計算進行優化並提供很大的int8算力，好比RK3399 pro搭載的NPU提供3T int8算力。

Why 8 bit works?

那麼爲何int8在DL模型中可以應用呢？不是有數值精度損失麼？主要緣由有兩個：

1. 訓練好的DNN網絡時出了名的對噪聲和擾動魯棒性強。
2. 大部分訓練好的權重都落在一個很小的區間內。

通常正常操做的話，8 bit量化只會形成很低的精度損失，並且這個損失是能夠經過微調重訓練進行彌補的。好比在Han等人在這篇文章裏對AlexNet其中一層的權重進行分析：
<imge="https://cdn-images-1.medium.com/max/1600/0*UkgbJuMdr6eOBjux.png" style="zoom:50%"/>

左圖是實際權重，大部分分佈在-0.1到0.1的範圍內，而右圖是進行了4bit量化後的權重數值分佈，4bit可以最大表示16個數值，所以大部分權重都有塌縮，可以保持原來的值的只有16個值。那麼若是進行8bit的量化，最大可以保持256個值，對原始權重的保留會更加完整，量化形成的數值損失會很小。

雖然目前已經有文章開始研究直接用低精度的數值進行訓練，好比這個，可是須要替換乘法操做，訓練很是複雜，並且效果也暫時不夠好，仍是處於實驗階段。這主要是目前的SGD等算法須要不斷積累小很是小的梯度值進行權重更新。

How 8 bit works？

那麼如何用int類型來表示float類型呢？最簡單的方式就是乘一個係數把float類型的小數部分轉換成整數部分，而後用這個轉換出來的整數進行計算，計算結果在還原成float。相似的，量化具體須要如下幾點：

1. 量化的變換必須是線性的，這樣才能確保計算結果可以映射會原始值
2. 量化必須可以保持0.f的精度，這是由於0在DNN中做用很大，若是原始的0映射到量化後的值變成了其餘值並出現了精度損失，那個在計算過程當中就會引入誤差。

所以對於實際值和量化值的映射關係，通常能夠用如下公式表示：

\(r= (r_{max}-r_{min})/(2^B-1)-0*(q-z)\)

其中，r表示實際值；q表示量化的比特數，好比int8量化就是8；z表示量化後的0點值。

具體的映射關係以下:

從公式中能夠看到，量化的重要一點就是要肯定合適的\(r_{max}\)和\(r_{min}\).對於訓練好的模型就行post-training 量化來講，這點比較容易，只須要統計凍存的全部權重參數便可。

Post training quantization

通常來講，凍好的模型中典型的conv層包含如下參數：

• weights tensor
• input tensor
• forward pass operator
• output tensor

對輸出來講，大部分層輸出的值都只會落在一個很窄的區間內，所以對output進行量化就須要利用在在訓練的時候統計大部分輸入獲得的輸出來進行統計肯定合適的最大和最小值。

可是對於operation來講，直接利用以前的量化方式須要注意一點：因爲計算包括乘法，所以有可能會有計算結果的值溢出量化值（好比int8）所能表達的範圍(overflow)。所以這裏經常使用的方法是首先將結果用較大數值空間的量化值（好比int32）進行存儲，而後再量化到8 bit。採用int32至少能夠徹底覆蓋原始float類型運算的結果而不擔憂溢出。

此外，對於一些操做，其邏輯須要改變。好比ReLU，量化後須要比較的值是quantize(0)而不是原來的0.f。

如上圖所示，權重，輸入首先進行量化，其中輸入的量化範圍是須要訓練數據做爲參考；而輸出進行兩步量化，來兼容計算過程可能出現的溢出。

Fake Quantization in TFLite

在TensorFlow中，量化是經過fake quantization node來進行的。對於大模型來講，冗餘參數比較多，直接量化的影響比較小；可是對於小模型來講，冗餘參數就比較少了，直接量化致使的 精度損失可能會比較大。在fake quantizaiton中，訓練太重就就會模擬評估量化帶來的round effect在inference的時候的影響，所以在訓練過程當中量化值仍是會議float類型保存，並能夠經過反向傳播進行調整。具體quantization aware training能夠查看這裏.

此外，就如以前所說的fake quantizaiton node會記錄計算和輸出值的範圍用於量化。

Result

下表記錄了量化形成的精度損失，整體來看仍是很是小的。

What's next

到這位置，只是介紹了8bit量化在TF上的實現原理和方式。而如今學術界對於量化有不少新的研究包括：quantized training, non-linear quantization, binary quantization, networks without multipliers等等，在不久的未來但願可以出現無損且高效的量化手段，這將極大收益訓練和推測。