两行代码自动压缩ViT模型！模型体积减小3.9倍，推理加速7.1倍

Transformer 模型及其变体，因其更优的注意力机制能力和长时依赖等特性，已成为自然语言处理（NLP）、语音识别（ASR）、计算机视觉（CV）等领域的主流序列建模结构。根据 PaperWithCode 网站 ObjectDetectiononCOCO-test-dev 专栏数据，检测任务中 Transformer 类模型精度早已超越卷积类模型。然而，Transformer 模型的体积和速度相比卷积网络却存在很大的劣势，限制了其在产业上的应用。

本文以 ViT 这一典型的 Transformer 模型为例，为大家介绍使用低成本、高收益的 AI 模型自动压缩工具（ACT，Auto Compression Toolkit）。ACT 无需修改训练源代码，调用两行压缩 API，通过几十分钟量化训练，保证模型精度的同时，极大的减小模型体积，降低显存占用，提升模型推理速度，助力 AI 模型的快速落地。使用 ACT 工具压缩 ViT 模型，通过几十分钟量化训练，模型体积即可减小3.93倍。在 GPU 上，INT8 推理速度相比 FP32 推理速度提升7.1倍。

注：ViT 模型和分割任务模型测速 batch_size=40，其他模型测速 batch_size=1

本文将从以下五个方面进行讲解，全文大约3600字，预计阅读时长5分钟。

• ViT 模型分析
• ViT 压缩策略介绍
• ViT 自动化压缩
• ViT 推理优化
• 未来工作展望

 ViT 模型分析 
本节首先对 ViT 模型结构进行简单介绍，然后对各个模块的计算量进行分析，以确定压缩的重点。
▎ViT 模型结构分析
ViT 使用标准 Transformer 结构替代传统的卷积操作，将图像拆分为多 patch 后再输入 Transformer 中。对于一张3x224x224大小的输出图片，首先会被分成14x14个 patch，每个 patch 的大小为3x16x16，每个 patch 都使用一个768x3x16x16的卷积核进行特征提取，生成1x1x768的特征图，整个图片经过卷积之后得到的特征大小为1x14x14x768，再对特征进行维度展开，得到1x196x768的特征。借鉴 NLP 场景，在每个图片的特征前引入一个 CLS 字段的特征，来帮助模型进行分类。最终得到完整的特征是1x197x768维度，197对应的是 NLP 场景中的每个句子的 token 数量，768对应的是 NLP 场景中每个 token 的维度。接下来这些特征就可以直接输入到后续堆叠的 Transformer Encoder 结构中。如下图所示：
注：图片来自 ViT 论文^[5]

Transformer Encoder 结构是由一个 MultiHead Self-Attention（简称 MSA）和一个 MLP 组成。MSA 中包含六个矩阵乘计算，其中四个全连接层对 MSA 不同部分的特征进行特征映射或特征融合，两个 Batch GEMM 和一个 softmax 算子表示特征的自注意力机制。MLP 是由一个 GeLU 非线性层和两个全连接线性层组成。整个 Transformer Encoder 结构如下图所示：

▎计算量分析
再来分析下 Transformer Encoder 每一个模块对应的计算量，方便后续对各个模块进行压缩。从上图可以看出在整个 Transformer Encoder 中，在整个模型中，大部分计算量在于矩阵乘计算，所以接下来着重分析每个模块的矩阵乘计算量。

• MSA 计算公式如下：

其中,表示 U_qkv 三个矩阵乘的参数，D 代表每个 patch 映射到的特征维度，qk^T 对应上图中 Batch GEMM1，N 表示图片有多少个 patch，Av 对应上图中的 Batch Gemm2，MSA(z)计算公式对应的是上图中的 Proj Gemm，k 表示 head 数量。

• MLP 计算公式如下：

MLP 在计算过程中对特征维度进行了扩展，利用 U₁ 和 U₂ 两个参数特征扩展，所以 MLP 中每个矩阵乘的计算量是 MSA 中每个矩阵乘的计算量的4倍。
每个矩阵乘在整个 Transformer Encoder 矩阵乘计算量占比如下图所示：
注：QKV GEMM、FC1 GEMM 以及 FC2 GEMM 的矩阵乘计算量占比总和达到89%，所以后续重点介绍如何压缩这些矩阵乘。
 ViT 压缩策略解析 
模型压缩中最常用的技术包括量化、剪枝、知识蒸馏和神经网络搜索。压缩 ViT 模型主要用到量化、剪枝和知识蒸馏，自动压缩技术主要使用这三种方法来对 ViT 模型进行整体压缩，在保证压缩模型精度的前提下，减小模型体积并带来实际的推理加速。
本章节简单的介绍这三个方法的原理，主要介绍如何把这些方法应用在 ViT 模型上，对技术原理感兴趣的同学可以查看之前的文章（见文末文章链接）。
▎量化
量化本质上就是用低比特数值存储计算代替高精度数值存储和计算，来达到节省显存和加速计算的目的。在本案例中，量化主要针对最开始的卷积、MSA 中 QKV 三个矩阵乘、Proj 矩阵乘和 MLP 中的两个矩阵乘。量化包括离线量化和量化训练两种方式。

• 方式一：离线量化

离线量化只需要一些校准数据，收集模型中间的激活值，用一些离线量化策略来处理收集到的激活值，得到量化映射参数。然而，对一些模型进行离线量化会造成一定的精度损失，因此，有一些研究专门针对 ViT 这种特殊结构的模型进行离线量化算法改进，例如利用 attention map 重要性机制来构造量化损失^[1]，或者利用 log 量化^[2]等方式来对 ViT 模型进行量化，可以提升量化后模型的精度或者把模型量化到更低比特。

• 方式二：量化训练

量化训练需要构建量化模型，使用训练数据微调模型，得到训练后的量化模型，这种方式相比离线量化通常会有更优的精度。有一些研究针对 ViT 结构来做量化训练，例如根据 MSA 的特性把量化参数和量化比特数都作为可学习的参数参与整体量化模型训练^[6]，把 ViT 模型量化到更合适的比特数。
▎剪枝
剪枝是指减掉模型的参数维度，直接减少了计算量，达到降低显存和加速的目的。剪枝包括非结构化剪枝、半结构化剪枝和结构化剪枝。其中，非结构化剪枝和半结构化剪枝需要推理库的特殊支持，例如针对1x1卷积的非结构化剪枝，PaddleSlim 联合 Paddle Lite 对其进行了特殊优化，在精度无损的情况下带来实际的加速。英伟达针对安培架构 ASP 稀疏也实现了特殊的推理。因此，对于非结构化剪枝和半结构化剪枝如果没有推理优化，很难有实际的加速。
为了保证压缩后的模型在任一硬件和任一推理库上都有实际的推理加速，建议选择结构化剪枝，减掉的参数量就是实际减少的计算量，无需推理库特殊的支持。
和量化类似，我们选择对 ViT 中计算量大的矩阵乘进行剪枝，可选的剪枝策略也比较多，这里主要介绍结构化剪枝，例如对 ViT 模型构造剪枝空间^[3]，对剪枝空间中可调整的部分进行重要性排序，迭代式的进行结构化剪枝，得到最终的剪枝模型。
▎蒸馏
蒸馏主要是把精度较高、复杂度较高的模型的信息传递给精度较低、复杂度较低的模型，从而提升小模型的精度。相比大模型，蒸馏后的模型在保证模型精度的前提下降低了模型的复杂度。
最早针对 ViT 进行蒸馏的 DeIT^[4]在学生模型中，通过在刚开始引入一个可学习的蒸馏标识符来学习教师模型的输出，用蒸馏标识符和 ViT 模型本身引入的分类标识符同时来优化学生模型，达到蒸馏的目的。
 ViT 模型自动压缩流程 
上一章介绍了 ViT 模型进行压缩时主要用到的策略，本章节介绍如何使用 ACT 工具通过配置文件配置压缩策略，实现 ViT 模型的自动压缩。
▎准备待压缩模型
准备待压缩的部署模型。可以从 PaddleClas 套件模型库、自动压缩 repo 中获得待压缩的 ViT 模型，也可以自己训练好 ViT 模型。
▎准备数据读取模块
在对 ViT 模型进行自动化压缩的过程中，我们可以直接复用飞桨图像分类套件 PaddleClas 中相应的数据读取模块。

from ppcls.data import build_dataloader
data_config = config.get_config(global_config['reader_config'], show=False)
train_loader = build_dataloader(data_config["DataLoader"], "Train", device)

▎定义策略配置文件
>> 策略自动组合
自动化压缩工具提供的压缩策略有很多，可以在配置文件中指定。自动压缩工具会对这些配置进行自动组合，得到最终的组合压缩策略。以下内容会介绍如何在配置文件中指定压缩策略。

• 策略组合一：量化蒸馏训练配置

本节使用 ACT 中的量化训练和蒸馏训练对 ViT 进行量化，量化位置主要为卷积层、MSA 中的全连接层和 MLP 中的全连接层，减少模型显存的同时加速矩阵乘的计算速度。
在进行量化训练或者剪枝训练的时候需要有损失优化信息来反向传播优化模型参数，而 ACT 主要压缩的推理模型是不包含原始的损失优化信息，所以 ACT 中选择使用蒸馏来监督压缩后模型进行优化训练，用压缩前模型的隐藏层特征来监督压缩后模型的隐藏层特征，从而实现模型的训练。
ACT 量化配置如下：

1Distillation:                #蒸馏
2node: softmax_12.tmp_0
3
4QuantAware:                  #量化
5use_pact: true
6onnx_format: true
7quantize_op_types: [conv2d, matmul_v2]

*参数含义如下：Distillation 表示使用蒸馏；node 参数表示使用的蒸馏结点，这里选择的 ViT 网络结构最后输出的 softmax 输出结点作为蒸馏节点；QuantAware 表示使用量化训练的压缩策略，量化 op 选择 conv2d 和 matmul_v2 即可表示量化整个 ViT 模型中的卷积和全连接层；use_pact 为是否使用 PACT 量化；onnx_format 量化后的模型如果想转为 ONNX 格式，需要设置 onnx_format 为 true。

• 策略组合二：策略剪枝蒸馏训练

ACT 提供了 Transformer 通道数剪枝，主要是对 MSA 中的 head 数量、全连接层和 MLP 中的全连接层进行参数裁剪。剪枝逻辑为在剪枝前对模型进行 head 重要性排序和参数重要性排序，然后根据传入的剪枝比例对相应位置进行裁剪，得到剪枝后的模型。同样的，剪枝训练也需要结合蒸馏使用。ViT 模型主体结构是 Transformer，所以也可以使用 ACT 提供的 Transformer 通道数剪枝方案进行剪枝（目前还在适配中，可能您读到这篇推文或者看到直播的时候已经可以使用啦）。可以先看下应该如何设置对应的 ACT 剪枝配置：

1Distillation:
2node: softmax_12.tmp_0
3TransformerPrune:
4pruned_ratio: 0.25

*参数含义如下：Distillation 表示使用蒸馏；node 参数表示使用的蒸馏结点；TransformerPrune 表示使用 Transformer 剪枝策略进行剪枝；pruned_ratio 为剪枝比例，本方案设为25%，您的可以根据需要自行设置；剪枝和量化压缩是两种交互使用的压缩策略，最优的方式是串行使用这两种压缩策略，当然也可以单独使用。
>> 策略自动选择
针对 ViT 模型，如果没有指定具体的压缩策略（即配置文件中没有指定策略），在剪枝已经支持的情况下会先对模型进行结构化剪枝，再对剪枝后的模型进行量化，从而得到最终优化后的模型。在剪枝没有支持的情况下会自动选择量化的策略进行模型压缩。
🔗完整 ViT 自动压缩量化压缩配置文件链接：https://github.com/PaddlePaddle/PaddleSlim/blob/develop/example/auto_compression/image_classification/configs/VIT/qat_dis.yaml
>> 两行核心自动压缩代码
启动 ACT 时，需要传入模型文件的路径（model_dir）、模型文件名（model_filename）、参数文件名称（params_filename）、压缩后模型存储路径（save_dir）、压缩配置文件（config）、dataloader 和评估精度的 eval_callback 参数。

1ac=AutoCompression(model_dir=global_config['model_dir'],
2                   model_filename=global_config['model_filename'],
3                   params_filename=global_config['params_filename'],
4                   save_dir=args.save_dir,
5                   config=all_config,
6                   train_dataloader=train_dataloader,
7                   eval_callback=eval_function if rank_id==0 else None,
8                   eval_dataloader=eval_dataloader)
9ac.compress()

 ViT 推理优化 
服务器推理引擎 Paddle Inference 团队对 ViT 模型的推理进行了一些特殊优化，来进一步提升压缩后模型的推理速度。主要包括以下四种：
▎算子融合
对于一些常用算子进行算子融合，减少整个模型算子的个数，加快模型的推理。算子融合规则如下图所示：

▎灵活使用 TensorRT 推理库
在输入图像较小的情况下，attention 部分的计算可以利用 TensorRT 高性能推理库，采用预置策略，在保证性能的同时兼容易用性。
▎基于硬件特性调优
在输入图像较大的情况下，attention 部分的计算需要进行一些其他优化，由于大部分硬件是对8的倍数的矩阵乘有着更优的计算性能，所以在模型计算过程中会把 attention 部分的计算逻辑 padding 到8的倍数来进一步提升矩阵乘的性能，引入的 padding 算子也和 transpose 算子进行了算子融合，在保证性能的前提下增加 attention 部分的灵活性。
▎实现高效计算算子
另外也根据显卡特性比如向量化，开发更加高效的计算 kernel 来加速整个 ViT 模型的推理。
为了更好地介绍飞桨自动化压缩技术在 ViT 模型的应用，解答近半年自动压缩社区开发者使用中遇到的高频问题，我们准备了本次直播课，快快扫码加群关注我们的直播间吧！
直播时间：2022.12.13（周二）20:30-21:30。

⭐Star 收藏⭐欢迎大家点 star 鼓励并前来体验https://github.com/PaddlePaddle/PaddleSlim/tree/develop/example/auto_compression
 未来工作展望 

• 飞桨自动化压缩工具 ACT 将支持更多 AI 模型的压缩，扩展支持的模型数量和应用场景。
• 将 ACT 的各种能力（包括模型剪枝、非结构化稀疏等压缩方法）在不同类型模型上进行验证。
• 升级 ACT 能力，增加更多前沿的压缩算法。
• 支持完善更多部署方法，包括 ONNX Runtime、OpenVINO 等，进一步助力 AI 模型工程落地。

 往期自动压缩专题 

▎项目地址

• GitHub

https://github.com/PaddlePaddle/PaddleSlim/tree/develop/example/auto_compression

• Gitee

https://gitee.com/PaddlePaddle/PaddleSlim/tree/develop/example/auto_compression
▎参考文献[1] Post-Training Quantization for Vision Transformer[2] FQ-ViT: Post-Training Quantization for Fully Quantized Vision Transformer[3] NViT: Vision Transformer Compression and Parameter Redistribution[4] Training data-efficient image transformers & distillation through attention[5] AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE[6] Q-ViT: Fully Differentiable Quantization for Vision Transformer