模型壓縮：模型量化打怪升級(jí)之路 - 工具篇

掐指一算，距上篇文章發(fā)出剛剛好過(guò)去了一個(gè)月。幾件事情提醒我是時(shí)候更新了。

哪幾件事情呢？

1. 最近發(fā)現(xiàn)一些還在學(xué)校讀書的同學(xué)非常關(guān)注一個(gè)量化工作精度的高低，讀過(guò)我上篇分享(龔睿昊：模型壓縮：模型量化打怪升級(jí)之路 - 0 序章 https://zhuanlan.zhihu.com/p/349820682)的同學(xué)應(yīng)該知道，部分學(xué)術(shù)界的工作與工業(yè)界的實(shí)用性是存在一定脫節(jié)的。我個(gè)人的看法是：除了部分比較本質(zhì)的探究量化提升的研究工作之外，一套好用的、通用的、強(qiáng)大的模型量化工具是更值得我們?nèi)ネ苿?dòng)的一件事情。
2. 恰逢P(guān)yTorch1.8.0發(fā)布，里面推出了一個(gè)非常驚喜又意料之中的模塊：torch.fx (torch.fx - PyTorch 1.8.0 documentation  https://pytorch.org/docs/stable/fx.html)。這個(gè)模塊是干嘛的？它可以symbolic的trace一個(gè)PyTorch Module的forward，得到對(duì)應(yīng)的graph，并支持對(duì)graph做一定的transformation，然后還支持一個(gè)python2python的code generation，生成一個(gè)新的forward代碼。大家知道動(dòng)態(tài)圖特性一直是PyTorch的一大優(yōu)勢(shì)，而也恰恰成了它的一個(gè)短板：很難對(duì)計(jì)算圖做一些變換和分析。因此相繼有JIT、Torchscript和fx推出，有了這些我們可以干什么呢？一個(gè)很重要的方面是它使一些compilation相關(guān)的優(yōu)化有了更多空間，此外還有一個(gè)我自己最最最最關(guān)心的是，它可以更好的方便我們"自動(dòng)"插入量化節(jié)點(diǎn)。試圖在PyTorch上做過(guò) 量化自動(dòng)插入 的同學(xué)，應(yīng)該知道這里面的痛。不得不說(shuō)，這是一個(gè)好消息，但是有了它之后是否完備了呢？

目前看還是不夠的。今天我們就來(lái)針對(duì)模型量化的工業(yè)級(jí)生產(chǎn)框架/工具進(jìn)行一個(gè)大致的review，嘗試整理和總結(jié)已有框架的特點(diǎn)和能力，結(jié)合一些落地經(jīng)驗(yàn)提出我們認(rèn)為的下一代量化框架應(yīng)該具備的特性。

現(xiàn)有框架的特點(diǎn)和能力

開源的涉及模型量化的項(xiàng)目主要分為工業(yè)界主導(dǎo)、學(xué)術(shù)界主導(dǎo)和有志個(gè)人主導(dǎo)。

其中工業(yè)界會(huì)做量化工具并開源的主要有三類：深度學(xué)習(xí)框架廠商(Facebook/Tensorflow等)、硬件廠商(NVIDIA/Intel/Qualcomm等)以及一些致力于推進(jìn)相關(guān)工具作為云服務(wù)附加值的廠商(Microsoft等)，這些廠商所開源的工具與實(shí)際落地部署的貼合性受到其開發(fā)背景的影響，一般來(lái)說(shuō)框架廠商和硬件廠商的項(xiàng)目是更加符合量化的實(shí)際情況的（也不絕對(duì)）。學(xué)術(shù)界開源的項(xiàng)目主要是一些從事相關(guān)領(lǐng)域的研究員開發(fā)，這些項(xiàng)目主要以復(fù)現(xiàn)學(xué)術(shù)算法為主，一般不能保證能在硬件上實(shí)際部署。此外，也有一些其他的民間研究者開發(fā)的相關(guān)項(xiàng)目，其中有一些是很實(shí)際的。

本文將會(huì)選擇筆者所知道并看過(guò)的項(xiàng)目中的一部分重點(diǎn)進(jìn)行分析。其他的歡迎大家一起分享和補(bǔ)充～

• PyTorch：：

PyTorch最早期在量化方面是非常空白的，自1.4.x開始，引入了非常初步的量化支持，定義了一些low level API，比如：fuse_modules之類的功能函數(shù)用于做一些量化相關(guān)的算子合并(實(shí)際上非常粗糙，是不考慮forward執(zhí)行順序的fuse)；QuantStub/Observer/FakeQuantize之類的類方便用戶拼接出自己的quantized model。為了降低大家的使用復(fù)雜度，甚至也推出了一套model zoo的quantization版本。可以明顯看到的是，這個(gè)時(shí)期的量化支持是非常手工的，還遠(yuǎn)沒有在graph層面做量化替換的工作。

由于那時(shí)候，我們已經(jīng)做了一個(gè)叫"_Dirichlet"_的框架來(lái)做pytorch quantization的自動(dòng)插入這件事情了，所以只是覺得PyTorch的API設(shè)計(jì)是很簡(jiǎn)潔清晰的，但距離更低成本的使用還有很大距離，因此沒有過(guò)多的follow。而是針對(duì)Dirichlet本身仍存在的問(wèn)題思考如何進(jìn)一步改進(jìn)和提升。

隨著PyTorch推出了fx這一套，graph mode quantization的概念浮出水面。PyTorch官方也情理之中的對(duì)相關(guān)缺陷進(jìn)行了補(bǔ)齊，看完fx __init__.py（https://github.com/pytorch/pytorch/blob/master/torch/fx/__init__.py）中對(duì)其功能的大概注釋之后，我的心情非常復(fù)雜。萬(wàn)萬(wàn)沒想到PyTorch官方會(huì)用這種方式來(lái)解決問(wèn)題，因?yàn)榇蟾?年前在我們自己的框架Dirichlet中就引入了類似的機(jī)制來(lái)做自動(dòng)替換，這算是對(duì)我們選取的技術(shù)路線的一個(gè)官方肯定。但同時(shí)，在Dirichlet大量的業(yè)務(wù)落地檢驗(yàn)中，發(fā)現(xiàn)這種方式還是有一些挑戰(zhàn)需要解決的。因此對(duì)PyTorch引入這種機(jī)制之后的長(zhǎng)遠(yuǎn)規(guī)劃有了更進(jìn)一步的好奇。

首先簡(jiǎn)單介紹一下我們的做法：

上圖是我們的代碼替換機(jī)制，實(shí)際上這一步主要是為了做一些elementwise的提取和替換，看過(guò)torch.fx文檔的同學(xué)應(yīng)該發(fā)現(xiàn)了，跟其中的樣例非常像。此外我們同時(shí)還有一個(gè)借助PyTorch JIT trace進(jìn)行圖匹配和變換的模塊，兩者配合完成模型的自動(dòng)量化插入。其實(shí)從思想上跟當(dāng)前的PyTorch官方思路已經(jīng)非常一致了，不過(guò)我們主要是針對(duì)量化的一些場(chǎng)景進(jìn)行了設(shè)計(jì)，沒有像當(dāng)前的fx一樣非常嚴(yán)格的抽象成一個(gè)單獨(dú)的模塊。可以看到PyTorch現(xiàn)在抽象出一個(gè)單獨(dú)的fx模塊，應(yīng)該不僅僅是考慮量化的，能夠預(yù)想到其未來(lái)應(yīng)該會(huì)針對(duì)廣義的模型壓縮技術(shù)(pruning/nas)擴(kuò)展更多的支持。加上PyTorch團(tuán)隊(duì)專注于框架開發(fā)，因此在穩(wěn)定性以及與框架的協(xié)同發(fā)展上肯定是更有保證的。不過(guò)即便如此，基于Dirichlet的大量業(yè)務(wù)實(shí)踐，我們認(rèn)為仍有幾個(gè)挑戰(zhàn)需要PyTorch在未來(lái)解決：

1. 這樣一個(gè)code parse和generation的能力是有邊界的，如何定義清楚這個(gè)邊界，并讓用戶習(xí)慣和接受，是有一定難度的。舉一個(gè)量化的例子，如果在一個(gè)寫法比較放飛自我的模型定義上，自動(dòng)量化插入崩了，用戶其實(shí)是很難受的，他并不知道怎么改，尤其是在面臨不同的量化硬件量化插入模式可能還不一樣的情況下。
2. 一些特殊的模型定義不容易保證插入語(yǔ)義的正確性。比如目標(biāo)檢測(cè)里經(jīng)常出現(xiàn)的share head。即head中的conv權(quán)重是共享的，會(huì)在多個(gè)分支中forward多次，但是由于處理的輸入不同，他們的量化參數(shù)應(yīng)該是每個(gè)分支獨(dú)享的。shared head定義一般是for loop形式，如fx文檔所說(shuō)，這是目前所不支持的。
3. 對(duì)于分布式的擴(kuò)展性和支持。

雖然還有很多挑戰(zhàn)，但是PyTorch的發(fā)力思路是清晰的，只不過(guò)沒有經(jīng)過(guò)大量業(yè)務(wù)場(chǎng)景的打磨，很多細(xì)節(jié)需要補(bǔ)齊，更多的量化硬件支持需要補(bǔ)齊，不過(guò)整體來(lái)說(shuō)我對(duì)PyTorch量化的未來(lái)是持樂觀態(tài)度的。

• TensorFlow

TensorFlow算是框架層面支持量化的鼻祖了，從whitepaper（https://arxiv.org/abs/1806.08342）開始，Google算是把量化推向工業(yè)級(jí)實(shí)用的關(guān)鍵角色，其靜態(tài)圖特性決定了框架層面的量化支持沒有太多阻礙。加上Tensorflow Lite的推理支持，是端到端打通的典范了。當(dāng)然他的特點(diǎn)也是主要針對(duì)自家的推理實(shí)現(xiàn)進(jìn)行支持，不過(guò)很多硬件廠商會(huì)主動(dòng)支持tf格式的量化模型的導(dǎo)入，這就是另一回事了。

• NNCF

NNCF是intel推出的專門針對(duì)OpenVINO的模型壓縮框架，也有一篇專門的tech report（https://arxiv.org/abs/2002.08679）對(duì)其進(jìn)行介紹。個(gè)人認(rèn)為這是目前看到的最多的考慮了實(shí)用性并經(jīng)過(guò)了較多打磨的一個(gè)庫(kù)，不足之處是他忽略了一個(gè)非常重要的特性，不支持QAT中的merge BN。大概說(shuō)一下NNCF的特點(diǎn)：

1. 模仿PyTorch JIT自行擴(kuò)展了一套trace機(jī)制，用于控制context和靈活的量化插入。
2. 量化模塊設(shè)計(jì)為節(jié)點(diǎn)模式，便于靈活的量化插入和導(dǎo)出。
3. 維護(hù)了compression controller和context，甚至能夠處理shared head問(wèn)題。
4. 能夠?qū)С鰩uantize Node的onnx模型，進(jìn)一步部署到自家的OpenVINO上跑。
5. 有硬件量化模式定義的json配置，理論上可擴(kuò)展硬件種類支持和配置量化位置。
6. 一些模擬量化操作使用了cuda優(yōu)化。

其中很多是我們之前也做過(guò)的，比如自己實(shí)現(xiàn)一套trace，實(shí)現(xiàn)一些模擬量化的cuda extension，所以在PyTorch fx系列出來(lái)之前，PyTorch上做量化，我們是最服NNCF的。當(dāng)然它也有一些問(wèn)題，比如擴(kuò)展OP支持需要熟悉源代碼，不支持QAT中的merge BN，僅重點(diǎn)支持OpenVINO的量化推理。量化算法部分，只集成了hawq和一些learning based方法。

• AIMET

AIMET是高通推出的模型壓縮工具，對(duì)應(yīng)的推理庫(kù)是SNPE。AIMET使用的量化節(jié)點(diǎn)插入方式是module替換，這一點(diǎn)從設(shè)計(jì)先進(jìn)性上來(lái)說(shuō)，是弱于NNCF的。其中支持了一些自家研究員提出的算法如：Bias Correction，Weight Equalization，AdaRound等。

• TensorRT-PyTorch-Quantization

TensorRT官方提出的在PyTorch層面做量化支持的工具，支持量化參數(shù)的導(dǎo)出以寫入TensorRT engine中。不禁想起了當(dāng)年的TensorRT對(duì)齊經(jīng)歷，做PyTorch模擬量化在不確定很多硬件底層實(shí)現(xiàn)細(xì)節(jié)的情況下，是非常非常難做到對(duì)齊的。包括后來(lái)的NNIE量化對(duì)齊，讓我養(yǎng)成了對(duì)齊虐我千百遍，我待對(duì)齊如初戀的良好工具人意識(shí)。這里拋出一個(gè)小問(wèn)題，如下的計(jì)算圖局部，TensorRT是如何插入量化節(jié)點(diǎn)的？歡迎知道的在評(píng)論區(qū)留言～

• Distiller

Distiller是比較早開源的模型壓縮項(xiàng)目，由Intel Habana實(shí)驗(yàn)室開發(fā)。其最早相對(duì)是沒有那么貼近部署的，不過(guò)發(fā)展過(guò)程中逐漸擴(kuò)展了merge BN，quantize onnx export支持，還貼心的添加了與PyTorch quantize onnx的相互轉(zhuǎn)化。其中支持了DoReFa、PACT、自家研究員發(fā)明的LAPQ等算法。整體上相比兄弟單位的NNCF在落地和設(shè)計(jì)上略遜一籌。

• Vitis / Brevitis

VITIS-AI是Xilinx針對(duì)自家FPGA的一套模型壓縮推理工具套件，Brevitis是QAT模塊。不難看出，這套工具其實(shí)是深鑒的班子被Xilinx收購(gòu)后持續(xù)推進(jìn)的結(jié)果。其中的量化工具部分在去年5月份還比較粗糙，不支持量化參數(shù)的外部寫入，去年中旬發(fā)布的新版本(不記得確切的版本號(hào)了)開始首次支持。主要是針對(duì)自家hardware，去年5月份我摸這套工具的時(shí)候，可以猜出其量化方案采用的是比較老舊的Ristretto方案，不知道后續(xù)有沒有升級(jí)。

• TVM

TVM是由以陳天奇大佬發(fā)起的深度學(xué)習(xí)編譯器項(xiàng)目，其中也有int8 甚至int4/2/1 kernel的編譯支持。我們也做過(guò)大規(guī)模測(cè)試，其中一些kernel的性能表現(xiàn)還是不錯(cuò)的。這里重點(diǎn)想講的是其中的量化算法支持。可以看到TVM支持的離線量化算法是比較少的，主要有KL divergence的校準(zhǔn)。QAT算法目前主要靠其他框架的量化模型導(dǎo)入。個(gè)人認(rèn)為這一塊其實(shí)是可以進(jìn)一步提升的，很早就看到社區(qū)ziheng對(duì)更加自動(dòng)化的離線量化的RFC，前段時(shí)間終于看到了PR和對(duì)應(yīng)的tech report（https://github.com/apache/tvm）。最近也有UC berkeley的HAWQ-V3使用TVM進(jìn)行了測(cè)試。我們自己也在持續(xù)跟進(jìn)TVM社區(qū)的進(jìn)展，并已經(jīng)貢獻(xiàn)了一些小修小補(bǔ)的PR，劇透一下后續(xù)會(huì)有更大的一些PR提出來(lái)，希望為社區(qū)貢獻(xiàn)一份力量。

• PaddleSlim

依托PaddlePaddle的模型壓縮算法庫(kù)，其中有一些基于Paddle生態(tài)的算法復(fù)現(xiàn)，也注意了典型硬件比如TensorRT的對(duì)齊和導(dǎo)出。

• Condensa

也是NV推出的，但是非常不符合實(shí)際量化部署的一套庫(kù)。這就可以看出，像NV這樣的核彈廠，確實(shí)有很多團(tuán)隊(duì)在重復(fù)造輪子，抬手就可以列出一堆：NVIDIA/sampleQAT、Post Training Quantization (PTQ)，但是目前沒有看到比較系統(tǒng)完整又出色的出來(lái)。

NVIDIA/sampleQAT：https://github.com/NVIDIA/sampleQAT

Post Training Quantization (PTQ)：https://nvidia.github.io/TRTorch/tutorials/ptq.html

下一代量化框架還有多長(zhǎng)的路要走？

上面只是列舉出了一些工業(yè)界推出的框架，可以看到大有百花齊放的態(tài)勢(shì)，針對(duì)PyTorch的量化，其中從設(shè)計(jì)來(lái)說(shuō)個(gè)人最為看好的是PyTorch官方的方案和NNCF的方案。當(dāng)然其他有幾個(gè)也有很大的發(fā)展空間。遺憾的是，所有這些框架和工具都只能解決一部分問(wèn)題。我們距離下一代量化框架，到底還有多長(zhǎng)的路要走？個(gè)人看來(lái)，一個(gè)優(yōu)秀的量化框架應(yīng)該具備以下幾個(gè)特點(diǎn)：

1. 同時(shí)具有針對(duì)developer的low level API和針對(duì)小白用戶的high level API。
2. 靈活、魯棒、可靠的量化插入和替換機(jī)制。
3. 可擴(kuò)展、可配置的多種量化硬件支持，包括NNIE/DSP/GPU/CPU/ASIC等等等等。
4. 強(qiáng)大的PTQ和QAT算法支持，以及靈活的量化pipeline控制機(jī)制，最大化算法潛力。
5. 大量的應(yīng)用場(chǎng)景模型量化實(shí)踐打磨，包括多種任務(wù)和多種模型。

其中的1、2已經(jīng)通過(guò)PyTorch官方和NNCF衍生出了兩種技術(shù)方案，未來(lái)大概率會(huì)逐漸收斂。而3、4、5是已有框架解決的不夠好的，本質(zhì)上的原因還是他們離應(yīng)用場(chǎng)景還不夠近，只會(huì)聚焦一部分需求。這方面我們其實(shí)是有天然優(yōu)勢(shì)的。因此我們也會(huì)一方面吸收開源社區(qū)的優(yōu)秀設(shè)計(jì)模式，同時(shí)整理我們自己的落地和算法經(jīng)驗(yàn)，推動(dòng)下一代量化框架的開花和結(jié)果。