PTQ4SAM模型论文解读，并基于MindSpore NLP推理复现

作者：linyiining

来源：开源实习

Segment Anything Model（SAM）在许多计算机视觉任务中取得了令人印象深刻的性能。然而，作为一种大规模模型，其巨大的内存和计算成本限制了其实际部署。该论文针对SAM量化中固有的瓶颈问题，提出了一种针对SAM的后训练量化（PTQ）框架，命名为PTQ4SAM。该框架在各种视觉任务（实例分割、语义分割和目标检测）、数据集和模型变体上的广泛实验结果表明了PTQ4SAM的优越性。

# 01 论文创新点（SAM模型量化的固有挑战）

1、post-key-linear的双峰分布现象影响了模型量化后的表现

在 SAM（Segment Anything Model） 这样的视觉大模型中，Post-Key-Linear 这一术语通常与 Transformer 架构中的注意力机制（Attention Mechanism） 相关。具体来说，它指的是在计算注意力机制中的 Key 矩阵 后，经过一个线性变换（Linear Layer）的输出。SAM模型的Post-Key-Linear输出呈现双峰分布（如下图1(a)所示），这极大影响了量化的效果。

解决方法：双峰集成策略。文章通过per-tensor和per-channel两个角度分析了双峰分布，并认为这是量化的主要障碍，并将双峰分布转换为正态分布来解决这一瓶颈。

**2、**复杂的 post Softmax 分布

由于注意力机制的种类多种多样，与 VIT 相比，SAM 表现出更复杂的 post Softmax 分布。文章中举了image-to-token和token-to-image两个位置的post-Softmax的不同分布说明这一点，并且指出先前的工作并未意识到这一点，导致了固有信息的潜在丢失。

解决方法：自适应粒度量化（AGQ）

# 02 PTQ4SAM架构

1、双峰集成策略

从per-tensor角度来看激活值双峰分布，如图1(a)所示。从per-channel角度来看，每个通道的激活值在一个固定值附近，每个通道之间有很大区别，如图2所示。

由于每一个channel中所有tensor的激活值都分布在一个固定值附近，而这一系列固定值呈现以0为中心的双峰分布。因此，论文中采用一个参数γ来将双峰分布转换为正常分布，其取值为-1或1。对于channel的激活值平均值为正的，γ为1，channel激活值不发生变化；对于平均值为负的，γ为-1，将负值映射为正值，同时为query linear也乘上-1，保持输出值不变。由于query linear是关于0对称的标准正态分布，乘以-1后分布不变。这样，q和k的linear激活值都保持了标准分布。

代码的实现主要分为三步骤：①判断tensor是否属于双峰分布②计算γ③进行等效转换（bimodal discovery, γ computation and equivalent transformation）。

2、自适应粒度量化（AGQ）

对post softmax的量化，用的是改进的log2量化，一种自适应粒度量化方法。通过搜索最优的2的幂次基底τ，为不同的Softmax后分布提供合适的量化粒度。通过建立一个LUT（对不同的τ和移位的乘积进行查找），就可以通过非常小的空间代价（4KB）实现对于不同post softmax的不同粒度的量化。这个空间代价是用所有的τ的取值乘以所有的移位可能。对于8bit的精度来说，τ的取值有2^2种可能性，aq有2^n（n取8）种可能，-aq/τ总共需要4byte来存（包含第一部分的aq%τ和第二部分的[aq/τ]，所以一共是4KB。

移位操作：通过移位来快速实现2的幂次方的乘除法计算。

选定τ的方法：通常来说是直接衡量attention矩阵A的损失，文中采用了AV的乘积损失作为目标函数，这样能够直接衡量在模型整体性能上的量化损失，单独衡量A的损失和整个注意力块的损失并不一致。通过在Calibration Set上实验，对不同的softmax层选择不同的τ数值。同时论文发现，τ取值较小的时候能够对低注意力分数更好的量化，而τ增大后，对更大的注意力分数的量化性能更好。

3、预训练编码器和解码器权重共享的，**这种方法不仅能提高模型性能，还能显著降低模型的内存占用。共享权重的方式在大多数任务中都表现出了良好的效果，尤其在计算资源有限的情况下尤为重要。

# 03 实验

1、实验设置

任务和数据集：实验涵盖了实例分割（MS-COCO数据集）、语义分割（ADE20K数据集）和目标检测（DOTA-v1.0数据集）等任务。

模型：使用了SAM-B、SAM-L和SAM-H三种模型变体。

量化设置：实验中使用了不同的量化比特宽度（W6A6和W4A4），分别表示权重和激活的比特数。

2、实验结果

实例分割：在MS-COCO数据集上，PTQ4SAM方法在不同的检测器中始终优于其他方法。同时，我们的 PTQ4SAML 令人鼓舞地实现了无损精度，如下图3，在 W6A6 设置下，我们的 PTQ4SAM-L 在 SAM-L 上应用 YOLOX 和 H-Deformable-DETR 时分别达到 40.3% 和 41.2%，与全精度型号相比，性能仅下降 0.1% 和 0.3%。

语义分割：在ADE20K数据集上，PTQ4SAM在W6A6量化时甚至超过了全精度模型的性能。如下图4，SAM-L模型在W6A6量化时达到了33.66%的mIOU，比全精度模型高出0.05%。

目标检测：在DOTA-v1.0数据集上D 定向目标检测中，我们的方法始终优于其他基于学习的 PTQ 方法。PTQ4SAM在SAM-L和SAM-H上进行W6A6量化时仅比全精度模型下降了0.3%，而在W4A4时，仍然能够保持较高的性能。

3、主****要贡献

1）第一个针对SAM的PTQ策略。

2)提出了针对双峰分布的双峰集成策略。

3)提出了AGQ方法。

4)通过实验证明，我们的方法是一个即插即用的PTQ方法，大大超过了之前最先进的方法。

# 04 MindSpore推理验证

首先，复现论文代码得到sam_b的量化模型。

然后，通过mindspore加载COCO数据集，下面是实例分割数据集的定义，每次迭代会返回"img_id", "image_path", "image", "bboxes", "labels", "masks"。文件命名dataset.py。（此处仅展示部分代码，点击阅读原文可查看完整代码）

import os

定义推理代码。为了实现实例分割，需要加载yolox作为检测头，然后把yolox检测出的boxes作为SAM模型的提示，输出目标掩码。首先，借助mindyolo库实现yolox目标检测mindyolo_test.py：

import math

然后实现推理代码。为了评估，输出需要规范化为COCO格式，存入json中。

import torch

调用COCO工具评估准确率：

import os

测试结果：

可视化展示

# 05 总结

PTQ4SAM是首个专门针对SAM模型的后训练量化解决方案，通过BIG和AGQ策略，有效解决了SAM模型量化中的双峰分布和多样的Softmax后分布问题。在多种视觉任务和不同模型变体上，PTQ4SAM均表现出优越的性能，能够在低比特量化下保持较高的精度，并实现显著的计算和存储节省。

尽管PTQ4SAM在量化SAM模型方面取得了显著成果，但论文也指出，SAM模型中双峰分布的成因尚不清楚，这将是未来研究的一个潜在方向。