TPU-MLIR量化敏感層分析，提升模型推理精度

背景介紹

TPU-MLIR編譯器可以將機(jī)器學(xué)習(xí)模型轉(zhuǎn)換成算能芯片上運行的bmodel模型。由于浮點數(shù)的計算需要消耗更多的計算資源和存儲空間，實際應(yīng)用中往往采用量化后的模型（也稱定點模型）進(jìn)行推理。相比于浮點數(shù)模型，量化模型的推理精度會有一定程度的損失。當(dāng)精度損失較大時，需要搜索模型中對精度影響較大的層，即敏感層，將其改回浮點類型，生成混精度模型進(jìn)行推理。

以mobilenet-v2網(wǎng)絡(luò)為例，使用ILSVRC-2012數(shù)據(jù)驗證集的5萬張圖片驗證浮點數(shù)模型和量化模型（表格中分別記為FLOAT和INT8）的精度，INT8模型的Top1精度降低了3.2%，Top5精度降低了2%。

敏感層搜索

TPU-MLIR的敏感層搜索功能會計算網(wǎng)絡(luò)模型中的每一層分別由浮點數(shù)類型轉(zhuǎn)成定點數(shù)類型后，對網(wǎng)絡(luò)模型輸出造成的損失。同時，由于量化threshold值也會影響定點模型的精度，敏感層搜索功能考慮了三種量化方法——KL、MAX和Percentile對精度的影響。KL方法首先統(tǒng)計FLOAT模型tensor絕對值的直方圖（2048個bin），得到參考概率分布P，隨后用INT8類型去模擬表達(dá)這個直方圖(128個bin)，得到量化概率分布Q，在不同的截取位置計算P和Q的KL散度，最小散度對應(yīng)的截取位置記為KL方法得到的threshold值。MAX方法是用FLOAT模型tensor絕對值的最大值作為量化threshold。Percentile方法則是通過統(tǒng)計FLOAT模型tensor絕對值的百分位數(shù)來確定threshold。

算法流程

敏感層搜索算法的流程圖如下：

敏感層搜索流程圖

1. 在搜索之前需要先判斷FLOAT模型和INT8模型的相似度，用一定數(shù)目的圖片（比如30張）進(jìn)行推理，計算兩種模型輸出結(jié)果的余弦相似度的平均值，如果平均值達(dá)到預(yù)設(shè)值，比如0.99，則認(rèn)為INT8模型與FLOAT模型的相似度較高，不需要進(jìn)行敏感層搜索。
2. 對于該模型，生成三種量化方法對應(yīng)的量化表。
3. 循環(huán)每個op及每種量化方法，將該op改為INT8類型，并采用對應(yīng)量化方法下的threshold，生成混精度模型，計算其與FLOAT模型網(wǎng)絡(luò)輸出的loss（1減余弦相似度），記錄最優(yōu)threshold，即最低loss值對應(yīng)的threshold，隨后將該op改回FLOAT類型。
4. 將所有op按照loss從大到小排序，選擇排名前五的op生成qtable。

在搜索過程中，對于每個op及其每種threshold值下生成的混精度模型與FLOAT模型的loss，以及最終按照loss排序的所有op信息（包括op的名字、類型和loss），均會記錄在log日志中，用戶可以通過查看log日志，手動調(diào)整qtable文件。

使用方法

敏感層搜索需要輸入從model_transform步驟得到的mlir文件，run_calibration步驟得到的量化表，推理用的數(shù)據(jù)集等，使用命令如下：

run_sensitive_layer.pymobilenet.mlir\
--dataset../ILSVRC2012\
--input_num100\
--inference_num30\
--max_float_layers5\
--expected_cos0.99\
--post_processpostprocess.py\
--calibration_tablemobilenet_cali_table\
--chipbm1684\
-omobilenet_qtable

各項參數(shù)的含義如下表所示：

敏感層搜索程序會生成如下文件：

1. 用于生成混精度模型的qtable，每行記錄了需要轉(zhuǎn)回FLOAT的op及轉(zhuǎn)換類型，例如：input3.1 F32；
2. 經(jīng)過調(diào)優(yōu)后的新量化表new_cali_table，在原始量化表的基礎(chǔ)上，更新了每個op的threshold值為三種量化方法中最優(yōu)的threshold；
3. 搜索日志SensitiveLayerSearch，記錄整個搜索過程中，每個op在每種量化方法下，混精度模型與FLOAT模型的loss;

注意，在model_deploy步驟生成混精度模型時，需要使用qtable和新量化表。

精度測試結(jié)果

仍以前述精度測試采用的mobilenet-v2網(wǎng)絡(luò)為例，使用ILSVRC2012數(shù)據(jù)集中的100張圖片做量化，30張圖片做推理，敏感層搜索總共耗時402秒，占用內(nèi)存800M，輸出結(jié)果信息如下：

the layer input3.1 is 0 sensitive layer, loss is 0.008808857469573828, type is top.Conv
the layer input11.1 is 1 sensitive layer, loss is 0.0016958347875666302, type is top.Conv
the layer input128.1 is 2 sensitive layer, loss is 0.0015641432811860367, type is top.Conv
the layer input130.1 is 3 sensitive layer, loss is 0.0014325751094084183, type is top.Scale
the layer input127.1 is 4 sensitive layer, loss is 0.0011817314259702227, type is top.Add
the layer input13.1 is 5 sensitive layer, loss is 0.001018420214596527, type is top.Scale
the layer 787 is 6 sensitive layer, loss is 0.0008603856180608993, type is top.Scale
the layer input2.1 is 7 sensitive layer, loss is 0.0007558935451825732, type is top.Scale
the layer input119.1 is 8 sensitive layer, loss is 0.000727441637624282, type is top.Add
the layer input0.1 is 9 sensitive layer, loss is 0.0007138056757098887, type is top.Conv
the layer input110.1 is 10 sensitive layer, loss is 0.000662179506136229, type is top.Conv
......
run result:
int8 outputs_cos:0.978847 old
mix model outputs_cos:0.989741
Output mix quantization table to mobilenet_qtable
total time:402.15848112106323

觀察可知，input3.1的loss最大，且值為其他op的至少5倍。嘗試只將input3.1加進(jìn)qtable，其他層都保持INT8類型不變，生成混精度模型，在ILSVRC2012驗證集上進(jìn)行推理，精度如下：

上表中，MIX(oricali)代表使用原始量化表的混精度模型，MIX(newcali)代表使用新量化表的混精度模型，可以看出，基于三種量化方法的threshold調(diào)優(yōu)，也對模型的精度起到了正向影響。相比于INT8模型，混精度模型的Top1精度提升1.5%，Top5精度提升約1%。

對比混精度搜索

混精度搜索是TPU-MLIR中的另一個量化調(diào)優(yōu)搜索功能，它的核心思想是先尋找到layer_cos不滿足要求的層，再將該層及其下一層均從INT8轉(zhuǎn)回FLOAT，生成混精度模型，計算其與FLOAT模型輸出的余弦相似度，如果余弦相似度達(dá)到預(yù)設(shè)值，則停止搜索。注意，混精度搜索只有在op表現(xiàn)較差，即混精度模型與FLOAT模型輸出相似度低于INT8模型與FLOAT模型輸出相似度時，才會將op類型從FLOAT置回INT8，否則不會改變op的精度類型。所以混精度搜索不需要從頭進(jìn)行推理，耗時較短。這兩種方法的對比如下：

混精度搜索從網(wǎng)絡(luò)輸入開始，不斷搜索對網(wǎng)絡(luò)輸出相似度有提升的層加入qtable，直到網(wǎng)絡(luò)輸出相似度達(dá)到預(yù)設(shè)值后終止搜索。該方法可能會漏掉靠近網(wǎng)絡(luò)輸出部分的敏感層，而敏感層搜索方法會遍歷所有op，不會出現(xiàn)這種遺漏，這也是敏感層搜索的優(yōu)勢之處。在實際應(yīng)用中，用戶可以先使用速度較快的混精度搜索，如果沒能達(dá)到預(yù)期效果，再使用敏感層搜索功能進(jìn)行全局遍歷。

結(jié)語

敏感層搜索功能旨在尋找模型量化時對精度影響較大的層，它會遍歷模型中的所有op及三種量化方法，選擇最優(yōu)的量化threshold，記錄所有op的loss。將搜索到的敏感層設(shè)定為浮點數(shù)類型，其余層設(shè)定為定點數(shù)類型，生成混合精度模型，可以提升模型推理的精度。目前的敏感層搜索功能在mobilenet-v2網(wǎng)絡(luò)中表現(xiàn)優(yōu)異，只需要將loss最大的一層置為FLOAT，就可以獲得1.5%的精度提升。與混精度搜索方法相比，敏感層搜索雖然耗時更久，但它考慮了所有op和三種量化方法，不會遺漏靠近網(wǎng)絡(luò)輸出部分的敏感層。未來可以考慮從三個方面對敏感層搜索功能進(jìn)行優(yōu)化：1) 結(jié)合混精度搜索的優(yōu)點，在搜索時考慮每層及其鄰近層對網(wǎng)絡(luò)輸出的綜合影響；2) 在生成qtable時，不是根據(jù)用戶設(shè)定的數(shù)目，選擇loss最大的前N層，而是通過計算，將能讓網(wǎng)絡(luò)輸出相似度達(dá)到預(yù)設(shè)值的op均加入qtable；3) 在遍歷op的過程中考慮并行，縮短搜索時間。