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

背景介紹

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

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

敏感層搜索

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

算法流程

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

敏感層搜索流程圖

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

在搜索過程中，對(duì)于每個(gè)op及其每種threshold值下生成的混精度模型與FLOAT模型的loss，以及最終按照loss排序的所有op信息（包括op的名字、類型和loss），均會(huì)記錄在log日志中，用戶可以通過查看log日志，手動(dòng)調(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

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

敏感層搜索程序會(huì)生成如下文件：

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

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

精度測試結(jié)果

仍以前述精度測試采用的mobilenet-v2網(wǎng)絡(luò)為例，使用ILSVRC2012數(shù)據(jù)集中的100張圖片做量化，30張圖片做推理，敏感層搜索總共耗時(shí)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驗(yàn)證集上進(jìn)行推理，精度如下：

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

對(duì)比混精度搜索

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

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

結(jié)語

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