花生_TL007

这个屌丝很懒，什么也没留下！

热门标签

模型压缩：模型量化打怪升级之路-工具篇

作者：花生_TL007 | 2024-03-22 17:16:21

踩

shared head

本文转载自商汤泰坦公开课。

最近发现一些还在学校读书的同学非常关注一个量化工作精度的高低，读过我上篇分享(模型压缩：模型量化打怪升级之路 - 0 序章)的同学应该知道，部分学术界的工作与工业界的实用性是存在一定脱节的。

我个人的看法是：除了部分比较本质的探究量化提升的研究工作之外，一套好用的、通用的、强大的模型量化工具是更值得我们去推动的一件事情。

恰逢 PyTorch1.8.0 发布，里面推出了一个非常惊喜又意料之中的模块：torch.fx (torch.fx - PyTorch 1.8.0 documentation)。

torch.fx - PyTorch 1.8.0 documentation:

https://pytorch.org/docs/stable/fx.html

这个模块是干嘛的？它可以 symbolic 的 trace 一个 PyTorch Module的forward，得到对应的 graph，并支持对 graph 做一定的transformation，然后还支持一个 python2python 的 code generation，生成一个新的 forward 代码。

大家知道动态图特性一直是 PyTorch 的一大优势，而也恰恰成了它的一个短板：很难对计算图做一些变换和分析。因此相继有 JIT、Torchscript 和 fx 推出，有了这些我们可以干什么呢？一个很重要的方面是它使一些 compilation 相关的优化有了更多空间，此外还有一个我自己最最最最关心的是，它可以更好的方便我们"自动"插入量化节点。试图在 PyTorch 上做过量化自动插入的同学，应该知道这里面的痛。

不得不说，这是一个好消息，但是有了它之后是否完备了呢？

目前看还是不够的。今天我们就来针对模型量化的工业级生产框架/工具进行一个大致的 review，尝试整理和总结已有框架的特点和能力，结合一些落地经验提出我们认为的下一代量化框架应该具备的特性。

现有框架的特点和能力

开源的涉及模型量化的项目主要分为工业界主导、学术界主导和有志个人主导。

其中工业界会做量化工具并开源的主要有三类：深度学习框架厂商( Facebook/Tensorflow 等)、硬件厂商( NVIDIA/Intel/Qualcomm 等)以及一些致力于推进相关工具作为云服务附加值的厂商( Microsoft 等)，这些厂商所开源的工具与实际落地部署的贴合性受到其开发背景的影响，一般来说框架厂商和硬件厂商的项目是更加符合量化的实际情况的（也不绝对）。

学术界开源的项目主要是一些从事相关领域的研究员开发，这些项目主要以复现学术算法为主，一般不能保证能在硬件上实际部署。此外，也有一些其他的民间研究者开发的相关项目，其中有一些是很实际的。

本文将会选择笔者所知道并看过的项目中的一部分重点进行分析。其他的欢迎大家一起分享和补充～

PyTorch

PyTorch 最早期在量化方面是非常空白的，自1.4.x开始，引入了非常初步的量化支持，定义了一些 low level API，比如：fuse_modules 之类的功能函数用于做一些量化相关的算子合并(实际上非常粗糙，是不考虑 forward 执行顺序的 fuse)；QuantStub/Observer/FakeQuantize 之类的类方便用户拼接出自己的 quantized model。

PyTorch:

https://pytorch.org/docs/stable/quantization.html

为了降低大家的使用复杂度，甚至也推出了一套 model zoo 的 quantization 版本。可以明显看到的是，这个时期的量化支持是非常手工的，还远没有在 graph 层面做量化替换的工作。

由于那时候，我们已经做了一个叫" Dirichlet "的框架来做 pytorch quantization 的自动插入这件事情了，所以只是觉得 PyTorch 的 API 设计是很简洁清晰的，但距离更低成本的使用还有很大距离，因此没有过多的 follow。而是针对 Dirichlet 本身仍存在的问题思考如何进一步改进和提升。

随着 PyTorch 推出了 fx 这一套，graph mode quantization 的概念浮出水面。PyTorch 官方也情理之中的对相关缺陷进行了补齐，看完 fx __init__.py 中对其功能的大概注释之后，我的心情非常复杂。万万没想到 PyTorch 官方会用这种方式来解决问题，因为大概2年前在我们自己的框架 Dirichlet 中就引入了类似的机制来做自动替换，这算是对我们选取的技术路线的一个官方肯定。

fx __init__.py:

https://github.com/pytorch/pytorch/blob/master/torch/fx/__init__.py

但同时，在 Dirichlet 大量的业务落地检验中，发现这种方式还是有一些挑战需要解决的。因此对 PyTorch 引入这种机制之后的长远规划有了更进一步的好奇。

首先简单介绍一下我们的做法：

Dirichlet中的code_generation机制

上图是我们的代码替换机制，实际上这一步主要是为了做一些 elementwise 的提取和替换，看过 torch.fx 文档的同学应该发现了，跟其中的样例非常像。此外我们同时还有一个借助 PyTorch JIT trace 进行图匹配和变换的模块，两者配合完成模型的自动量化插入。

其实从思想上跟当前的 PyTorch 官方思路已经非常一致了，不过我们主要是针对量化的一些场景进行了设计，没有像当前的 fx 一样非常严格的抽象成一个单独的模块。可以看到 PyTorch 现在抽象出一个单独的fx模块，应该不仅仅是考虑量化的，能够预想到其未来应该会针对广义的模型压缩技术(pruning/nas)扩展更多的支持。加上 PyTorch 团队专注于框架开发，因此在稳定性以及与框架的协同发展上肯定是更有保证的。

不过即便如此，基于 Dirichlet 的大量业务实践，我们认为仍有几个挑战需要 PyTorch 在未来解决：

这样一个 code parse 和 generation 的能力是有边界的，如何定义清楚这个边界，并让用户习惯和接受，是有一定难度的。举一个量化的例子，如果在一个写法比较放飞自我的模型定义上，自动量化插入崩了，用户其实是很难受的，他并不知道怎么改，尤其是在面临不同的量化硬件量化插入模式可能还不一样的情况下。
一些特殊的模型定义不容易保证插入语义的正确性。比如目标检测里经常出现的 share head。即 head 中的 conv 权重是共享的，会在多个分支中 forward 多次，但是由于处理的输入不同，他们的量化参数应该是每个分支独享的。shared head 定义一般是 for loop 形式，如 fx 文档所说，这是目前所不支持的。
对于分布式的扩展性和支持。

目标检测中的share head举例：weight共享，输入的量化参数各自独享

虽然还有很多挑战，但是 PyTorch 的发力思路是清晰的，只不过没有经过大量业务场景的打磨，很多细节需要补齐，更多的量化硬件支持需要补齐，不过整体来说我对PyTorch量化的未来是持乐观态度的。

TensorFlow

TensorFlow 算是框架层面支持量化的鼻祖了，从 whitepaper 开始，Google 算是把量化推向工业级实用的关键角色，其静态图特性决定了框架层面的量化支持没有太多阻碍。加上Tensorflow Lite 的推理支持，是端到端打通的典范了。当然他的特点也是主要针对自家的推理实现进行支持，不过很多硬件厂商会主动支持tf格式的量化模型的导入，这就是另一回事了。

TensorFlow: https://github.com/tensorflow/model-optimization

whitepaper: https://arxiv.org/abs/1806.08342

NNCF

NNCF 是 intel 推出的专门针对 OpenVINO 的模型压缩框架，也有一篇专门的 tech report 对其进行介绍。个人认为这是目前看到的最多的考虑了实用性并经过了较多打磨的一个库，不足之处是他忽略了一个非常重要的特性，不支持 QAT 中的 merge BN。

NNCF: https://github.com/openvinotoolkit/nncf

tech report : https://arxiv.org/abs/2002.08679

大概说一下 NNCF 的特点：

模仿 PyTorch JIT 自行扩展了一套 trace 机制，用于控制 context 和灵活的量化插入。
量化模块设计为节点模式，便于灵活的量化插入和导出。
维护了 compression controller 和 context，甚至能够处理 shared head 问题。
能够导出带 Quantize Node 的 onnx 模型，进一步部署到自家的 OpenVINO 上跑。
有硬件量化模式定义的 json 配置，理论上可扩展硬件种类支持和配置量化位置。
一些模拟量化操作使用了 cuda 优化。

其中很多是我们之前也做过的，比如自己实现一套 trace，实现一些模拟量化的 cuda extension，所以在 PyTorch fx 系列出来之前，PyTorch 上做量化，我们是最服 NNCF 的。当然它也有一些问题，比如扩展 OP 支持需要熟悉源代码，不支持 QAT 中的merge BN，仅重点支持 OpenVINO 的量化推理。量化算法部分，只集成了 hawq 和一些 learning based 方法。

AIMET

AIMET 是高通推出的模型压缩工具，对应的推理库是 SNPE。AIMET 使用的量化节点插入方式是 module 替换，这一点从设计先进性上来说，是弱于 NNCF 的。其中支持了一些自家研究员提出的算法如：Bias Correction，Weight Equalization，AdaRound 等。

AIMET: https://github.com/quic/aimet

TensorRT-PyTorch-Quantization

TensorRT 官方提出的在 PyTorch 层面做量化支持的工具，支持量化参数的导出以写入 TensorRT engine 中。不禁想起了当年的 TensorRT 对齐经历，做 PyTorch 模拟量化在不确定很多硬件底层实现细节的情况下，是非常非常难做到对齐的。

TensorRT-PyTorch-Quantization:

https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization

包括后来的 NNIE 量化对齐，让我养成了对齐虐我千百遍，我待对齐如初恋的良好工具人意识。这里抛出一个小问题，如下的计算图局部，TensorRT 是如何插入量化节点的？欢迎知道的在评论区留言～

这个局部拓扑TensorRT会如何插入量化？

Distiller

Distiller 是比较早开源的模型压缩项目，由 Intel Habana 实验室开发。其最早相对是没有那么贴近部署的，不过发展过程中逐渐扩展了 merge BN，quantize onnx export 支持，还贴心的添加了与 PyTorch quantize onnx 的相互转化。其中支持了 DoReFa、PACT、自家研究员发明的 LAPQ 等算法。整体上相比兄弟单位的 NNCF 在落地和设计上略逊一筹。

Vitis / Brevitis

VITIS-AI 是 Xilinx 针对自家 FPGA 的一套模型压缩推理工具套件，Brevitis 是 QAT 模块。不难看出，这套工具其实是深鉴的班子被 Xilinx 收购后持续推进的结果。其中的量化工具部分在去年5月份还比较粗糙，不支持量化参数的外部写入，去年中旬发布的新版本(不记得确切的版本号了)开始首次支持。主要是针对自家 hardware，去年5月份我摸这套工具的时候，可以猜出其量化方案采用的是比较老旧的 Ristretto 方案，不知道后续有没有升级。

Vitis：https://github.com/Xilinx/Vitis-AI

Brevitis：https://github.com/Xilinx/brevitas

TVM

TVM 是由以陈天奇大佬发起的深度学习编译器项目，其中也有 int8 甚至 int4/2/1 kernel 的编译支持。我们也做过大规模测试，其中一些 kernel 的性能表现还是不错的。这里重点想讲的是其中的量化算法支持。可以看到 TVM 支持的离线量化算法是比较少的，主要有 KL divergence 的校准。QAT 算法目前主要靠其他框架的量化模型导入。个人认为这一块其实是可以进一步提升的，很早就看到社区 ziheng 对更加自动化的离线量化的 RFC，前段时间终于看到了 PR 和对应的 tech report。最近也有 UC berkeley 的 HAWQ-V3 使用 TVM 进行了测试。我们自己也在持续跟进 TVM 社区的进展，并已经贡献了一些小修小补的 PR，剧透一下后续会有更大的一些 PR 提出来，希望为社区贡献一份力量。

TVM：https://github.com/apache/tvm

tech report：https://github.com/apache/tvm

PaddleSlim

依托 PaddlePaddle 的模型压缩算法库，其中有一些基于 Paddle 生态的算法复现，也注意了典型硬件比如 TensorRT 的对齐和导出。

PaddleSlim：

https://github.com/PaddlePaddle/PaddleSlim

Condensa

也是 NV 推出的，但是非常不符合实际量化部署的一套库。这就可以看出，像 NV 这样的核弹厂，确实有很多团队在重复造轮子，抬手就可以列出一堆：NVIDIA/sampleQAT、Post Training Quantization (PTQ)，但是目前没有看到比较系统完整又出色的出来。

Condensa：https://github.com/NVlabs/condensa

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

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

下一代量化框架还有多长的路要走？

上面只是列举出了一些工业界推出的框架，可以看到大有百花齐放的态势，针对 PyTorch 的量化，其中从设计来说个人最为看好的是 PyTorch 官方的方案和 NNCF 的方案。当然其他有几个也有很大的发展空间。

遗憾的是，所有这些框架和工具都只能解决一部分问题。我们距离下一代量化框架，到底还有多长的路要走？

个人看来，一个优秀的量化框架应该具备以下几个特点：

同时具有针对 developer 的 low level API 和针对小白用户的 high level API。
灵活、鲁棒、可靠的量化插入和替换机制。
可扩展、可配置的多种量化硬件支持，包括 NNIE/DSP/GPU/CPU/ASIC 等等等等。
强大的 PTQ 和 QAT 算法支持，以及灵活的量化 pipeline 控制机制，最大化算法潜力。
大量的应用场景模型量化实践打磨，包括多种任务和多种模型。

其中的1、2已经通过 PyTorch 官方和 NNCF 衍生出了两种技术方案，未来大概率会逐渐收敛。而3、4、5是已有框架解决的不够好的，本质上的原因还是他们离应用场景还不够近，只会聚焦一部分需求。这方面我们其实是有天然优势的。因此我们也会一方面吸收开源社区的优秀设计模式，同时整理我们自己的落地和算法经验，推动下一代量化框架的开花和结果。

如果你耐心的读到了这里，对上述问题感兴趣并且有一定理解，有志跟我们一起造轮子推进下一代量化框架，欢迎直接私信或者把简历发至我的邮箱: gongruihao@sensetime.com，海量计算资源，业务实践场景等着你。

原文链接：

https://zhuanlan.zhihu.com/p/355598250

END