從TensorRT看INT8量化原理

本篇文章授權轉載於大神arleyzhang的《TensorRT(5)-INT8校准原理》https://arleyzhang.github.io/articles/923e2c40/，支持原創請查看原文。

另附GTC 2017，Szymon Migacz 的PPT 

Low Precision Inference

現有的深度學習框架 比如：TensorFlow，pytorch，Caffe， MixNet等，在訓練一個深度神經網絡時，往往都會使用 float 32（Full Precise ，簡稱FP32）的數據精度來表示，權值、偏置、激活值等。但是如果一個網絡很深的話，比如像VGG，ResNet這種，網絡參數是極其多的，計算量就更多了（比如VGG 19.6 billion FLOPS, ResNet-152 11.3 billion FLOPS）。如此多的計算量，如果中間值都使用 FP 32的精度來計算的話，勢必會很費時間。而這對於嵌入式設備或者移動設備來說，簡直就是噩夢，因為他們的計算能力和內存數量是不能與PC相比的。

因此解決此問題的方法之一就是在部署推理時（inference）使用低精度數據，比如INT8。除此之外，當然還有模型壓縮之類的方法，不過此處不做探究。注意此處只是針對 推理階段，訓練時仍然使用 FP32的精度。

從經驗上來分析一下低精度推理的可行性：

實際上有些人認為，即便在推理時使用低精度的數據（比如INT8），在提升速度的同時，也並不會造成太大的精度損失，比如 Why are Eight Bits Enough for Deep Neural Networks? 以及Low Precision Inference with TensorRT 這兩篇博文。

文章的作者認為網絡在訓練的過程中學習到了數據樣本的模式可分性，同時由於數據中存在的噪聲，使得網絡具有較強的魯棒性，也就是說在輸入樣本中做輕微的變動並不會過多的影響結果性能。與圖像上目標間的位置，姿態，角度等的變化程度相比，這些噪聲引進的變動只是很少的一部分，但實際上這些噪聲引進的變動同樣會使各個層的激活值輸出發生變動，然而卻對結果影響不大，也就是說訓練好的網絡對這些噪聲具有一定的容忍度（tolerance ）。

正是由於在訓練過程中使用高精度（FP32）的數值表示，才使得網絡具有一定的容忍度。訓練時使用高精度的數值表示，可以使得網絡以很小的計算量修正參數，這在網絡最后收斂的時候是很重要的，因為收斂的時候要求修正量很小很小（一般訓練初始 階段學習率稍大，越往后學習率越小）。

那么如果使用低精度的數據來表示網絡參數以及中間值的話，勢必會存在誤差，這個誤差某種程度上可以認為是一種噪聲。那也就是說，使用低精度數據引進的差異是在網絡的容忍度之內的，所以對結果不會產生太大影響。

以上分析都是基於經驗的，理論上的分析比較少，不過文章提到了兩篇 paper，如下：

• Improving the speed of neural networks on CPUs
• Training deep neural networks with low precision multiplications

這里不對這兩篇paper做探究。

TensorRT 的INT8模式只支持計算能力為6.1的GPU（Compute Capability 6.1 ），比如： GP102 (Tesla P40 and NVIDIA Titan X), GP104 (Tesla P4), and GP106 GPUs，主要根源是這些GPU支持 DP4A硬件指令。DP4A下面會稍微介紹一下。

TensorRT INT8 Inference

首先看一下不同精度的動態范圍:

	動態范圍	最小正數
FP32	−3.4×1038 +3.4×1038−3.4×1038 +3.4×1038	1.4×10−451.4×10−45
FP16	−65504 +65504−65504 +65504	5.96×10−85.96×10−8
INT8	−128 +127−128 +127	11

實際上將FP32的精度降為INT8還是比較具有挑戰性的。注：python的float類型式FP64的。

Quantization

將FP32降為INT8的過程相當於信息再編碼（re-encoding information ），就是原來使用32bit來表示一個tensor，現在使用8bit來表示一個tensor，還要求精度不能下降太多。

將FP32轉換為 INT8的操作需要針對每一層的輸入張量（tensor）和 網絡學習到的參數（learned parameters）進行。

首先能想到的最簡單的映射方式就是線性映射（或稱線性量化，linear quantization）, 就是說映射前后的關系滿足下式：

$FP32 \quad Tensor (T)=scale\_factor(sf) * 8-bit Tensor(t)+FP32\_bias (b)$

試驗證明，偏置實際上是不需要的，因此去掉偏置，也就是

$T=sf * t$

$sf$是每一層上每一個tensor的換算系數或稱比例因子（scaling factor），因此現在的問題就變成了如何確定比例因子。然后最簡單的方法是下圖這樣的:

• 簡單的將一個tensor 中的 -|max| 和 |max| FP32 value 映射為 -127 和 127 ，中間值按照線性關系進行映射。
• 稱這種映射關系為不飽和的（No saturation ），對稱的。

但是試驗結果顯示這樣做會導致比較大的精度損失。

下面這張圖展示的是不同網絡結構的不同layer的激活值分布，有卷積層，有池化層，他們之間的分布很不一樣，因此合理的 量化方式 應該適用於不同的激活值分布，並且減小 信息損失。因為從FP32到INT8其實就是一種信息再編碼的過程

筆者理解的直接使用線性量化的方式導致精度損失比較大的原因是：

• 上圖是一些網絡模型中間層的 激活值統計，橫坐標是激活值，縱坐標是統計數量的歸一化表示，這里是歸一化表示，不是絕對數值統計；
• 這個激活值統計 針對的是一批圖片，不同的圖片輸出的激活值不完全相同。所以圖上並不是一條曲線而是多條曲線（一張圖片對應一條曲線，或者稱為散點圖更好一點），只不過前面一部分重復在一塊了（紅色虛線圈起來的部分），說明對於不同圖片生成的大部分激活值其分布是相似的；但是在激活值比較大時（紅色實線圈起來的部分），曲線不重復了，一個激活值對應多個不同的統計量，這時的激活值分布就比較亂了。
• 后面這一部分在整個層中是占少數的（占比很小，比如10^-9, 10^-7, 10^-3），因此后面這一段完全可以不考慮到映射關系中去，保留激活值分布的主方向。開始我以為網絡之所以能把不同類別的圖片分開是由於后面實線部分的差異導致的，后來想了一下：這個並不包含空間位置的分布，只是數值上的分布，所以后面的應該對結果影響不大。

因此TensorRT的做法是:

• 這種做法不是將 ±|max| 映射為 ±127，而是存在一個 閾值 |T| ，將 ±|T| 映射為±127，顯然這里 |T|<|max|。
• 超出 閾值 ±|T| 外的直接映射為閾值 ±127。比如上圖中的三個紅色點，直接映射為-127。
• 稱這種映射關系為飽和的（Saturate ），不對稱的。
• 只要 閾值 選取得當，就能將分布散亂的較大的激活值舍棄掉，也就有可能使精度損失不至於降低太多。

網絡的前向計算涉及到兩部分數值：權值和激活值（weights 和activation，二者要做乘法運算），Szymon Migacz 也提到他們曾經做過實驗，說對weights 做saturation （飽和量化）沒有什么變化，因此 對於weights的int8量化就使用的是不飽和的方式；而對activation做saturation就有比較顯著的性能提升，因此對activation使用的是飽和的量化方式。

那現在的問題是 如何確定|T|？我們來思考一下，現在有一個FP32的tensor，FP32肯定是能夠表達這個tensor的最佳分布。現在我們要用一個不同的分布（INT8）來表達這個tensor，這個 INT8 分布不是一個最佳的分布。飽和的INT8分布由於閾值 |T|的取值會有很多種情況（128−|max|），其中肯定有一種情況是相對其他最接近FP32的，我們就是要把這種情況找出來。

既然如此，我們就需要一個衡量指標來衡量不同的 INT8 分布與原來的FP3F2分布之間的差異程度。這個衡量指標就是 相對熵（relative entropy），又稱為KL散度（Kullback–Leibler divergence，簡稱KLD），信息散度（information divergence），信息增益（information gain）。叫法實在太多了，最常見的就是相對熵。跟交叉熵也是有關系的。

• 假設我們要給一個信息進行完美編碼，那么最短平均編碼長度就是信息熵。

• 如果編碼方案不一定完美（由於對概率分布的估計不一定正確），這時的平均編碼長度就是交叉熵。

  平均編碼長度 = 最短平均編碼長度 + 一個增量

  交叉熵在深度學習中廣泛使用，衡量了測試集標簽分布和模型預測分布之間的差異程度。

• 編碼方法不一定完美時，平均編碼長度相對於最小值的增加量（即上面那個增量）是相對熵。

即 交叉熵=信息熵+相對熵

對於分類，因為是離散的，對於離散型隨機變量，信息熵公式如下：

$H(p)=H(X)=E_{x \sim p(x)} [-logp(x)]=- \sum_{i=1}^{n} p(x)log p(x)$

對於連續型隨機變量，信息熵公式如下：

$H(p)=H(X)=E_{x \sim p(x)}[-logp(x)]=- \int p(x)log p(x)dx$

通俗的理解 信息熵，交叉熵，相對熵，參考：知乎：如何通俗的解釋交叉熵與相對熵?

如何理解信息熵用來表示最短平均編碼長度，參考： 如何理解用信息熵來表示最短的平均編碼長度

詳細的不說了，請看參考鏈接。

在這里，FP32的tensor就是我們要表達的信息量，FP32也是最佳分布（可以認為最短編碼長度32bit），現在要做的是使用INT8 來編碼FP32的信息，同時要求INT8編碼后差異盡可能最小。考慮兩個分布 P（FP32）、Q（INT8）KL散度計算如下：

$D_{KL}(p||q)= \sum_{i=1}^{N} [p(x_{i})log p(x_{i})-p(x_{i})log q(x_{i})]= \sum_{i=1}^{N} p(x_{i})log \frac{p(x_{i})}{q(x_{i})}$

P，Q分別稱為 reference_distribution、 quantize _distribution

實際上這里也說明了每一層的tensor 的 |T| 值都是不一樣的。

確定每一層的 |T|值的過程稱為 校准（Calibration ）。

Calibration

上面已經說了 KL散度越小代表 INT8編碼后的信息損失越少。這一節來看看如何根據KL散度尋找最佳INT8分布。其實前面我們也已經提到了，如果要讓最后的精度損失不大，是要考慮一些先驗知識的，這個先驗知識就是每一層在 FP32精度下的激活值分布，只有根據這個才能找到更加合理的 閾值|T|。也就是說首先得有一個以FP32精度訓練好的模型。基本上現有的深度學習框架都是默認 FP32精度的，有些模型還支持FP16精度訓練，貌似 Caffe2和MXNet是支持FP16的，其他的不太清楚。所以基本上只要沒有特別設定，訓練出來的模型肯定是 FP32 的。

那激活值分布如何得到？難道我們要將FP32的模型先在所有的測試集（或驗證集）上跑一邊記錄下每一層的FP32激活值，然后再去推斷 |T|?

這里的做法是 從驗證集 選取一個子集作為校准集（Calibration Dataset ），校准集應該具有代表性，多樣性，最好是驗證集的一個子集，不應該只是分類類別的一小部分。激活值分布就是從校准集中得到的。

按照NVIDIA 官方的說法：

Note: The calibration set must be representative of the input provided to TensorRT at runtime; for example, for image classification networks, it should not consist of images from just a small subset of categories. For ImageNet networks, around 500 calibration images is adequate.

對於ImageNet 數據集來說 校准集大小一般500張圖片就夠了（Szymon Migacz的演講說用1000張），這里有點懷疑也有點震驚，沒想到 ImageNet 1000個分類，100多萬張圖片，500張就夠了，不過從2.5節的圖表中的結果可以看出500張確實夠了。

然后要做的是：

• 首先在 校准集上 進行 FP32 inference 推理；
• 對於網絡的每一層（遍歷）：
  • 收集這一層的激活值，並做 直方圖（histograms ），分成幾個組別（bins）（官方給的一個說明使用的是2048組），分組是為了下面遍歷 |T| 時，減少遍歷次數；
  • 對於不同的 閾值 |T| 進行遍歷，因為這里 |T|的取值肯定在 第128-2047 組之間，所以就選取每組的中間值進行遍歷；
    • 選取使得 KL_divergence(ref_distr, quant_distr) 取得最小值的 |T|。
• 返回一系列 |T|值，每一層都有一個 |T|。創建 CalibrationTable 。

上面解釋一下：假設 最后 使得 KL散度最小的|T|值是第200組的中間值，那么就把原來 第 0-200組的 數值線性映射到 0-128之間，超出范圍的直接映射到128

 校准的過程可以參考一下這個：https://www.jianshu.com/p/43318a3dc715 , 這篇文章提供了一個詳細的根據KL散度來將原始信息進行編碼的例子，包括直方圖的使用。跟這里的校准過程極為相像。

下面是一個官方 GTC2017 PPT 中給的校准的偽代碼：

//首先分成 2048個組，每組包含多個數值（基本都是小數）
Input: FP32 histogram H with 2048 bins: bin[ 0 ], …, bin[ 2047 ] 
  
For i in range( 128 , 2048 ): // |T|的取值肯定在 第128-2047 組之間,取每組的中點
    reference_distribution_P = [ bin[ 0 ] , ..., bin[ i-1 ] ] // 選取前 i 組構成P，i>=128
    outliers_count = sum( bin[ i ] , bin[ i+1 ] , … , bin[ 2047 ] ) //邊界外的組
    reference_distribution_P[ i-1 ] += outliers_count //邊界外的組加到邊界P[i-1]上，沒有直接丟掉
    P /= sum(P) // 歸一化
      
    // 將前面的P（包含i個組，i>=128），映射到 0-128 上，映射后的稱為Q，Q包含128個組，
    // 一個整數是一組
    candidate_distribution_Q = quantize [ bin[ 0 ], …, bin[ i-1 ] ] into 128 levels
    
    //這時的P（包含i個組，i>=128）和Q向量（包含128個組）的大小是不一樣的，無法直接計算二者的KL散度
    //因此需要將Q擴展為 i 個組，以保證跟P大小一樣
    expand candidate_distribution_Q to ‘ i ’ bins 
    
    Q /= sum(Q) // 歸一化
    //計算P和Q的KL散度
    divergence[ i ] = KL_divergence( reference_distribution_P, candidate_distribution_Q)
End For
//找出 divergence[ i ] 最小的數值，假設 divergence[m] 最小，
//那么|T|=( m + 0.5 ) * ( width of a bin )
Find index ‘m’ for which divergence[ m ] is minimal
threshold = ( m + 0.5 ) * ( width of a bin )

解釋一下第16行：

• 計算KL散度 KL_divergence(P, Q) 時，要求序列P和Q的長度一致，即 len(P) == len(Q)；
• Candidate_distribution_Q 是將 P 線性映射到 128個bins得到的，長度為128。而reference_distribution_P 包含 i （i>=128）個 bins （bin[0] - bin[i-1] ），二者長度不等；
• 需要將 candidate_distribution_Q 擴展回 i 個bins 然后才能與 i個bins 的 reference_distribution_P計算KL散度。

舉個簡單的栗子：

• 假設reference_distribution_P 包含 8 個bins（這里一個bin就只包含一個數據）:

  P = [ 1, 0, 2, 3, 5, 3, 1, 7]

• 我們想把它映射為 2 個bins，於是 4個一組合並：

  [1 + 0 + 2 + 3 , 5 + 3 + 1 + 7] = [6, 16]

• 然后要成比例的 擴展回到 8個組，保留原來是0的組：

  Q = [ 6/3, 0, 6/3, 6/3, 16/4, 16/4, 16/4, 16/4] = [ 2, 0, 2, 2, 4, 4, 4, 4]

• 然后對 P和Q進行標准化：

  P /= sum(P) 、Q /= sum(Q)

• 最后計算散度：

  result = KL_divergence(P, Q)

我們來看看 ResNet-152中 res4b30層校准前后的結果對比:

• 圖中那個白線就是 |T|的取值，不過怎么還小於128了，有點沒搞明白。

再看看其他幾種網絡的校准情況:

DP4A(Dot Product of 4 8-bits Accumulated to a 32-bit)

TensorRT 進行優化的方式是 DP4A (Dot Product of 4 8-bits Accumulated to a 32-bit)，如下圖:

這是PASCAL 系列GPU的硬件指令，INT8卷積就是使用這種方式進行的卷積計算。

這個沒搞太明白是怎么回事，參考這篇博客獲取詳細信息Mixed-Precision Programming with CUDA 8

下面是 官方 GTC2017 PPT 中給的INT8卷積計算的偽代碼：

// I8 input tensors: I8_input, I8_weights, INT8輸入tensor
// I8 output tensors: I8_output， INT8輸出tensor
// F32 bias (original bias from the F32 model),FP32的偏置
// F32 scaling factors: input_scale, output_scale, weights_scale[K], 這個是前面說的縮放因子sf
I32_gemm_out = I8_input * I8_weights // Compute INT8 GEMM (DP4A)，卷積計算，INT32輸出
F32_gemm_out = (float)I32_gemm_out // Cast I32 GEMM output to F32 float，強制轉換為FP32
//前面計算I8_input * I8_weights時，總的縮放系數為 input_scale * weights_scale[K]
//但是輸出的縮放系數為output_scale，所以為了保證縮放程度匹配，要將F32_gemm_out乘以 
//output_scale / (input_scale * weights_scale[ i ] )
  
// At this point we have F32_gemm_out which is scaled by ( input_scale * weights_scale[K] ),
// but to store the final result in int8 we need to have scale equal to "output_scale", so we have to rescale:
// (this multiplication is done in F32, *_gemm_out arrays are in NCHW format)
For i in 0, ... K-1:
rescaled_F32_gemm_out[ :, i, :, :] = F32_gemm_out[ :, i, :, :] * [ output_scale /(input_scale * weights_scale[ i ] ) ]
  
//將FP32精度的偏置 乘上縮放因子，加到前面的計算結果中
// Add bias, to perform addition we have to rescale original F32 bias so that it's scaled with "output_scale"
rescaled_F32_gemm_out _with_bias = rescaled_F32_gemm_out + output_scale * bias
//ReLU 激活
// Perform ReLU (in F32)
F32_result = ReLU(rescaled_F32_gemm_out _with_bias)
//重新轉換為 INT8
// Convert to INT8 and save to global
I8_output = Saturate( Round_to_nearest_integer( F32_result ) )

它這個INT8卷積的計算是這樣的，雖然輸入的tensor已經降為 INT8，但是在卷積計算的時候用了DP4A的計算模式，卷積計算完之后是INT32的，然后又要轉成 FP32，然后激活，最后再將FP32的轉為INT8.

只知道這么計算會快很多，但不知道為什么，詳情還是看Mixed-Precision Programming with CUDA 8 。

不過這個對於tensorRT的使用沒啥影響，這個是很底層的東西，涉及到硬件優化。

Typical workflow in TensorRT

典型的工作流還是直接使用 GTC2017 PPT 原文說法吧：

• You will need:
  • Model trained in FP32.
  • Calibration dataset.
• TensorRT will:
  • Run inference in FP32 on calibration dataset.
  • Collect required statistics.
  • Run calibration algorithm → optimal scaling factors.
  • Quantize FP32 weights → INT8.
  • Generate “CalibrationTable” and INT8 execution engine.

Results - Accuracy & Performance

精度並沒有損失太多

 速度提升還蠻多的，尤其是當 batch_size 大於1時，提升更明顯

 TITAN X GPU優化效果

 DRIVE PX 2, dGPU 優化效果

Open challenges / improvements

一些開放式的提升和挑戰：

• Unsigned int8 for activations after ReLU. 無符號 INT8 的映射。
• RNNs → open research problem. TensorRT 3.0開始已經支持RNN了。
• Fine tuning of saturation thresholds. 對閾值 |T|的 微調方法。
• Expose API for accepting custom, user provided scale factors. 開放API，使用戶可以自定義 換算系數（比例因子）

這幾個開放問題還是很值得研究的。

Conclusion

• 介紹了一種自動化，無參數的 FP32 到 INT8 的轉換方法；
• 對稱的，不飽和的線性量化，會導致精度損失較大；
• 通過最小化 KL散度來選擇 飽和量化中的 閾值 |T|;
• FP32完全可以降低為INT8推理，精度幾乎持平，速度有很大提升。

Reference

[1] https://arleyzhang.github.io/articles/923e2c40/