O_o >_< o_O O_o ~_~ o_O
大家好，我是極智視界，本文剖析一下 KL 對稱量化算法實現，以 Tengine 的實現為例。
前面已經寫過一篇《
1、KL 量化原理
KL 量化是用 KL 散度來衡量真實數據分布和量化數據分布之間的相似性的量化方法，是英偉達 TensorRT 中對於激活值采用的量化策略，KL 量化的主要邏輯如下：
-
KL 和 MIN-MAX 不一樣，不是直接將
[min, max]映射到[-127, 127]，而是去尋找一個閾值|T| < max(|max|, |min|)，將其[-T, T]映射到[-127, 127]。認為只要閾值選取得當，就能將閾值以外的值舍棄掉，也不會對精度損失造成大的影響； -
超出閾值
±|T|以外的值直接映射為閾值，如上圖中的三個紅色點，直接映射為 -127，這種映射關系稱為是飽和的。
KL 量化方法試圖將 float32 數值分布和 int8 數值分布抽象成兩個分布，用閾值 |T| 來更新這兩個數值分布，並用 KL 散度來衡量這兩個分布的相似性，若 KL 散度值越小，說明這兩個分布越相似，也就說明這個閾值 |T| 選擇的最好。對於對稱量化來說，根據這個閾值就能算出 Scale，而 Zero_point 始終為零。
下面的圖是 TensorRT 中的關於 KL 散度校准的偽代碼，這個圖也完美詮釋了 KLD 整個量化過程。(標記一下下圖為圖二，后面會調用)
2、KL 量化實現
這里還是以 Tengine 中 KL 量化的實現進行說明。
捋一下主要有以下幾個流程：
(1) 激活值量化：先求 min、max，再用 KL 策略搜索量化生成激活值校准表。fp32toint8；
(2) 權值量化：使用 min-max 量化策略。fp32toint8；
(3) 偏置量化：延用激活值量化 scale 進行 int32 量化。fp32toint32；
權值和偏置的量化比激活值量化多一步，除了要計算 Scale 外，還需要對值應用 Scale 進行直接量化以生成 int8 tmfile。
在 Tengine 中實現 KL 量化的主要代碼如下：
case ALGORITHM_KL:{
if (quant_tool.scale_file.empty()){
quant_tool.scale_file = "table_kl.scale";
quant_tool.activation_quant_tool();
}
save_graph_i8_perchannel(quant_tool.model_file.c_str(), quant_tool.scale_file.c_str(), quant_tool.output_file, quant_tool.inplace, false);
/* Evaluate quantitative losses */
if (quant_tool.evaluate){
fprintf(stderr, "[Quant Tools Info]: Step Evaluate, evaluate quantitative losses\n");
quant_tool.assess_quant_loss(0);
}
break;
}
其中最主要的量化搜索策略接口是 quant_tool.activation_quant_tool() 和 save_graph_i8_perchannel，對於 KL 量化來說這兩個接口分別做了兩件事：
(1) 激活值量化，生成 table_kl.scale；
(2) 權值&偏置量化，生成 scale_weight.txt、scale_bias.txt 和 int8 tmfile；
由於激活值量化中的 min、max 計算方式 及 權值&偏置量化過程，KL 量化和 MIN-MAX 量化邏輯相同且共用相同代碼，這里就不展開介紹了，這部分有興趣的同學可以查閱 《
KL 量化搜索策略的入口在這：
quant_tool.activation_quant_tool();
然后會先做 min、max 的比較搜索，主要用了 std::max_element、std::min_element 接口，這里不多說，得到 min、max 值后開啟 KL 搜索策略。
2.1 勾勒概率直方圖
做第一輪勾勒概率直方圖，進行第一輪的 KL 計算，第二輪開始不用重新勾勒概率直方圖，而是在第一輪構建的概率直方圖上進行迭代，所以你的校准圖片數量越多，這個最終得到的概率直方圖會越逼近真實分布。
/* calculate hist */
uint32_t inum = 0;
for (int i = 0; i < ir_graph->tensor_num; i++){
struct tensor* ir_tensor = ir_graph->tensor_list[i];
if (ir_tensor->tensor_type == TENSOR_TYPE_VAR || ir_tensor->tensor_type == TENSOR_TYPE_INPUT){
float step_max = std::abs(max_activation[i]);
if (std::abs(min_activation[i]) > step_max)
step_max = std::abs(min_activation[i]);
float step_bin = step_max / 2048.0f;
std::vector<float> every_edge;
if (nums == imgs_list.size() - 1){
for (int j = 0; j < 2048; j++){
float edge_float = (step_bin * (j + 0.5f));
every_edge.push_back(edge_float);
}
hist_edge.push_back(every_edge);
hist_gram.push_back(histCount((float*)ir_tensor->data, ir_tensor->elem_num, step_max));
}
else{
std::vector<uint32_t> hist_tmp;
hist_tmp = histCount((float*)ir_tensor->data, ir_tensor->elem_num, step_max);
for (int j = 0; j < 2048; j++){
hist_gram[inum][j] += hist_tmp[j];}
}
tensor_hist[i] = inum;
hist_tensor[inum] = i;
inum++;}
}
來看以下 histCount 接口：
std::vector<uint32_t> histCount(float* data, uint32_t elem_num, float abs_max){
float bin_scale = abs_max / 2047.f;
int bin_zp = 0;
std::vector<uint32_t> hist(2048);
for (int i = 0; i < elem_num; i++){
if (data[i] != 0){
uint32_t hist_idx = round(std::abs(data[i]) / bin_scale);
hist[hist_idx]++;}
}
return hist;
}
最后對得到的概率直方圖做一個歸一化處理：
distribution = normalize_histogram(distribution_in);
直方圖歸一化的實現接口也很簡單：
std::vector<float> normalize_histogram(std::vector<uint32_t>& histogram){
std::vector<float> histogram_out(histogram.size());
const size_t length = histogram.size();
float sum = 0;
for (size_t i = 1; i < length; i++)
sum += histogram[i];
for (size_t i = 1; i < length; i++)
histogram_out[i] = float(histogram[i] / sum);
return histogram_out;
}
2.2 計算 P
接下來的邏輯需要回頭看一下圖二，先計算 P 再計算 Q 最后計算 KL 散度。
先是計算模擬量化分布 P，從 target_bin = 128 --> 2048 遞增檢索，溢出部分映射到邊緣處理，可以把 P 認為是量化前 fp32 數據分布，即真實分布：
// get P
fill(quantize_distribution.begin(), quantize_distribution.end(), 0.0f);
const float num_per_bin = static_cast<float>(threshold) / static_cast<float>(target_bin);
for (int i = 0; i < target_bin; i++){
const float start = static_cast<float>(i) * num_per_bin;
const float end = start + num_per_bin;
const int left_upper = static_cast<int>(ceil(start));
if (static_cast<float>(left_upper) > start){
const float left_scale = static_cast<float>(left_upper) - start;
quantize_distribution[i] += left_scale * distribution[left_upper - 1];
}
const int right_lower = static_cast<int>(floor(end));
if (static_cast<float>(right_lower) < end){
const float right_scale = end - static_cast<float>(right_lower);
quantize_distribution[i] += right_scale * distribution[right_lower];
}
for (int j = left_upper; j < right_lower; j++){
quantize_distribution[i] += distribution[j];}
}
2.2 計算 Q
然后是計算真實量化分布 Q，伴隨 P 從 target_bin = 128 --> 2048 遞增檢索，可以把 Q 認為是量化后 int8 數據分布，即量化分布：
// get Q
std::vector<float> expand_distribution(threshold, 0);
for (int i = 0; i < target_bin; i++){
const float start = static_cast<float>(i) * num_per_bin;
const float end = start + num_per_bin;
float count = 0;
const int left_upper = static_cast<int>(ceil(start));
float left_scale = 0;
if (static_cast<float>(left_upper) > start){
left_scale = static_cast<float>(left_upper) - start;
if (distribution[left_upper - 1] != 0){
count += left_scale;}
}
const int right_lower = static_cast<int>(floor(end));
float right_scale = 0;
if (static_cast<float>(right_lower) < end){
right_scale = end - static_cast<float>(right_lower);
if (distribution[right_lower] != 0){
count += right_scale;}
}
for (int j = left_upper; j < right_lower; j++){
if (distribution[j] != 0){
count++;}
}
const float expand_value = quantize_distribution[i] / count;
if (static_cast<float>(left_upper) > start){
if (distribution[left_upper - 1] != 0){
expand_distribution[left_upper - 1] += expand_value * left_scale;}
}
if (static_cast<float>(right_lower) < end){
if (distribution[right_lower] != 0){
expand_distribution[right_lower] += expand_value * right_scale;}
}
for (int j = left_upper; j < right_lower; j++){
if (distribution[j] != 0){