Datawhale模型减肥秘籍Tasking之模型量化

什么是量化？为什么量化？

剪枝操作能够减少网络的参数量，从而起到压缩模型的作用。而量化 (quantization) 是另一种能够压缩模型参数的方法。量化将神经网络的浮点算法转换为定点，修改网络中每个参数占用的比特数，从而减少模型参数占用的空间。

移动端的硬件资源有限，比如内存和算力。而量化可以减少模型的大小，从而减少内存和算力的占用。同时，量化可以提高模型的推理速度。下图为不同数据类型的加法和乘法操作的耗时对比。

模型量化的好处如下：

• 减小模型大小：如 int8 量化可减少 75% 的模型大小，int8 量化模型大小一般为 32 位浮点模型大小的 1/4。
• 减少存储空间：在端侧存储空间不足时更具备意义。
• 减少内存占用：更小的模型当然就意味着不需要更多的内存空间。
• 减少设备功耗：内存耗用少了推理速度快了自然减少了设备功耗。
• 加快推理速度：浮点型可以访问四次 int8 整型，整型运算比浮点型运算更快；CPU 用 int8 计算的速度更快。

某些硬件加速器如 DSP/NPU 只支持 int8。比如有些微处理器属于 8 位的，低功耗运行浮点运算速度慢，需要进行 8bit 量化。

量化基本方法

这里可以理解成要进行模型压缩,就是利用尽可能少的权重表示全局的量,那么意思就是说把一个连续的序列取尽可能少的值来表示整体模型的取值分布,也就是区间离散化
主要两种方法

• 基于 k-means 的量化（K-means-based Quantization）：存储方式为整型权重 + 浮点型的转换表），计算方式为浮点计算。
• 线性量化（Linear Quantization）：存储方式为整型权重存储，计算方式为整型计算。

基于k-means的量化

k-means是一个经典的聚类方法,这里属于是将k-means运用到了模型压缩的内容当中,我认为这样子是因为k-means首先比较经典,再加上实现起来比较简单所以会比较常用
如下图所示，k-means 量化将weights聚类。每个权重的位置只需要存储聚类的索引值。将权重聚类成4类(0,1,2,3)，就可以实现2-bit的压缩。 存储占用从 32bit x 16 = 512 bit = 64 B => 2bit x 16 + 32 bit x 4 = 32 bit + 128 bit = 160 bit = 20 B
当weight更大时，压缩比例将会更大。
  
  下面是剪枝和k-means之间相结合的内容,首先循环得到最优的剪枝模型,然后利用量化的方法循环迭代尽可能少的权重
  

线性量化

线性量化是将原始浮点数据和量化后的定点数据之间建立一个简单的线性变换关系，因为卷积、全连接等网络层本身只是简单的线性计算，因此线性量化中可以直接用量化后的数据进行直接计算。

我们用 $r$表示浮点实数，$q$表示量化后的定点整数。浮点和整型之间的换算公式为：

$$r=S(q-Z)$$

$$q=round(r/S+Z)$$
其中，$S$是量化放缩的尺度，表示实数和整数之间的比例关系，$Z$是偏移量，表示浮点数中的 0 经过量化后对应的数（量化偏移），根据偏移量$Z$是否为0，可以将浮点数的线性量化分为对称量化（$Z=0$）和非对称量化（$Z\ne 0$）。大多数情况下量化是选用无符号整数，比如INT8的值域为[0,255],这种情况下需要要用非对称量化。
计算方法为：

$$S=\frac{r_{max}-r_{min}}{q_{max}-q_{min}}$$
$$Z=round(q_{max}-\frac{r_{max}}{S}$$
  其中，$r_{min}$ 和$r_{max}$分别表示浮点数中的最小值和最大值,$q_{min}$和$q_{max}$分别表示定点数中的最小值和最大值。

下面是计算的例子:

然后就可以计算量化后的矩阵:

这方面还有线性矩阵乘量化,最大值量化,全连接层线性量化等,其实基本上的思路就是给定一个规定,给一个范围的值映射到一个具体的值,从浮点数转成整数的方法去做

训练后量化

训练后量化（Post-Training Quantization, PTQ）是指在训练完成后，对模型进行量化，因此也叫做离线量化。

量化会带来精度损失，那么如何选取量化时所用参数（如scaling factor，zero point）可以尽可能地减少对准确率的影响呢？这也是我们需要关注的地方。量化误差来自两方面，一个是clip操作，一个是round操作。因此，我们还要介绍动态量化参数的计算方式，以及 round 这个操作带来的影响。

量化粒度

量化通常会导致模型精度下降。这就是量化粒度发挥作用的地方。选择正确的粒度有助于最大化量化，而不会大幅降低准确性性能。

逐张量量化（Per-Tensor Quantization）是指对每一层进行量化。在逐张量量化中，相同的量化参数应用于张量内的所有元素。在张量之间应用相同的参数会导致精度下降，因为张量内参数值的范围可能会有所不同。如下图的红框所示，3个channel共享一个量化参数。但是我们可以看到不同channel的数据范围是不同的。因此当 Layer-wise 量化效果不好时，需要对每个channel进行量化。

逐通道量化（Channel-wise Quantization就是将数据按照通道维度进行拆分，分别对每一通道的数据进行量化。相较于逐张量量化，逐通道量化可以减少量化误差，但需要更多的存储空间。逐通道量化可以更准确地捕获不同通道中的变化。这通常有助于 CNN 模型，因为不同通道的权重范围不同。由于现阶段模型越来越大，每个通道的参数也原来越多，参数的数值范围也越来越大，因此我们需要更细粒度的量化方式。

逐张量量化与逐通道量化的对比结果如下图所示。从图中可以看出：使用逐通道量化的误差更小，但付出的代价是必须存储更多信息(多个r和S) 。

组量化（Group Quantization）是指对通道内的数据拆分成多组向量，每组向量共享一个量化参数。VS-Quant 对张量的单个维度内的每个元素向量应用比例因子。它将通道维度细分为一组向量。

• 较小粒度时，使用较简单的整数缩放因子；
• 较大粒度时，使用较复杂的浮点缩放因子。
    存储开销：对于两级缩放因子，假设使用4-bit的量化，每16个元素有一个4-bit的向量缩放因子，那么有效位宽为 4+4/16=4.25bits。

为了提高能源效率，引入了两级缩放方案MX (Microscaling)。微缩放 (MX) 规范是从著名的 Microsoft 浮点 (MSFP) 数据类型升级而来的。该算法首先以每个向量的粒度计算浮点比例因子。然后，它通过将每向量比例因子分成整数逐向量分量和浮点逐通道分量来量化它们。MX 系列（如 MX4、MX6、MX9）表示了不同量化方案，它们的主要区别在于数据类型、缩放因子的设计以及组大小，目的在于通过压缩模型权重数据，优化神经网络的性能。

下图是不同的多级缩放方案对比结果。有效位宽 = (L0 数据位宽 + L0 量化尺度位宽 / L0 组大小 + L1 量化尺度位宽 / L1 组大小)。L0缩放因子通常采用较低的精度，用定点数表示，L1缩放因子则采用浮点数的表示方式。以MX6为例：L0数据类型是S1M4，表示1位符号位+4位尾数，共5位，用于表示数值。L0 量化尺度数据类型为E1M0，表示1位指数位，因此占1位。L0组大小是 2，这意味着L0量化尺度是针对每2个元素进行分组。L1 量化尺度数据类型为E8M0，表示8位指数位，因此占8 位。L1 组大小是 16，意味着 L1 量化尺度是针对每 16 个元素进行分组。所以可得到：有效位宽 = 5 + 1/2 + 8/16 = 6 bits。

动态量化参数的计算 ( Cliping )

指数移动平均（EMA）

指数移动平均（Exponential Moving Average, EMA）是一种常用的统计方法，用于计算数据的指数移动平均值。

EMA 收集了训练过程中激活函数的取值范围$r_{min}$和$r_{max}$，然后在每个 epoch 对这些取值范围进行平滑处理。

EMA的计算公式如下：

$$r_{\text{min, }{\max }^{t+1}}=\alpha r_{\text{min,max }}^t+(1-\alpha )r_{\text{min, }{}^{tax}}^{t+1}$$
其中，$r_{min,max}^t$表示第 t 步的取值范围,$\alpha$ 表示平滑系数。

Min-Max

Min-Max 是一种常用的校准方法，通过在训练好的 fp32 模型上跑少量的校准数据。统计校准数据的 $r_{min,max}$并取平均值作为量化参数。

KL 量化

KL 量化是用 KL 散度来衡量数据和量化后的数据之间的相似性；这种方法是去寻找一个阈值 $|T|<\max (|\max |,|\min |)$，将
$[-T,T]$映射到 [−127,128] 。并假设只要阈值选取得当，使得两个数据之间的分布相似，就不会对精度损失造成影响。

$$D_{KL}(P\Vert Q)=\sum _{i=1}^{n}P\left(x_i\right)\log \frac{P\left(x_i\right)}{Q\left(x_i\right)}$$

均方误差（MSE）

均方误差量化是指通过最小化输入数据 X 和量化后的数据 Q(X) 之间的均方误差，计算得到最合适的量化参数。

$$\underset{|r{|}_{\max }}{\min }E\left|(X-Q(X){)}^2\right|$$
  通过动态调整 $｜r｜max$来最小化均方误差。

Rounding

Rounding 是指将浮点数进行舍入操作，将浮点数映射到整数。最常用的 Rounding 方法是最近整数（Rounding-to-nearest）。权重是互相关联的，对每个权重的最好舍入不一定是对整个张量的最好舍入。如下图所示，如果我们考虑整体的数据分布，将权重 0.5 舍入为 1 不是一个好的选择。

我们最终想要的量化效果是输出数据的损失尽可能小，因此我们可以通过评判 rounding 对输出的影响来决定权重的舍入方式，也就是 AdaRound。简化的计算公式如下所示：

$$\mathrm{argmin}\Vert (Wx-\hat{W}x)\Vert$$
其中，

$\hat{W}=\lfloor \lfloor W\rfloor +\sigma \rceil ,\sigma \in [0,1]$，表示当前值是向上取整还是向下取整。

量化感知训练

量化感知训练（Quantization-Aware Training, QAT）是指在训练过程中，对模型添加模拟量化算子，模拟量化模型在推理阶段的舍入和裁剪操作，引入量化误差。并通过反向传播更新模型参数，使得模型在量化后和量化前保持一致。

前向传播

反向传播

实践内容

进行了kmeans量化的训练,训练函数和验证函数如下

def train(model: nn.Module,dataloader: DataLoader,criterion: nn.Module,optimizer: Optimizer,scheduler: LambdaLR,callbacks = None
) -> None:model.train()for inputs, targets in tqdm(dataloader, desc='train', leave=False):# Move the data from CPU to GPU# inputs = inputs.to('mps')# targets = targets.to('mps')# Reset the gradients (from the last iteration)optimizer.zero_grad()# Forward inferenceoutputs = model(inputs)loss = criterion(outputs, targets)# Backward propagationloss.backward()# Update optimizer and LR scheduleroptimizer.step()scheduler.step()if callbacks is not None:for callback in callbacks:callback()

@torch.inference_mode()
def evaluate(model: nn.Module,dataloader: DataLoader,extra_preprocess = None
) -> float:model.eval()num_samples = 0num_correct = 0for inputs, targets in tqdm(dataloader, desc="eval", leave=False):# Move the data from CPU to GPU# inputs = inputs.to('mps')if extra_preprocess is not None:for preprocess in extra_preprocess:inputs = preprocess(inputs)# targets = targets.to('mps')# Inferenceoutputs = model(inputs)# Convert logits to class indicesoutputs = outputs.argmax(dim=1)# Update metricsnum_samples += targets.size(0)num_correct += (outputs == targets).sum()return (num_correct / num_samples * 100).item()

构建k-means量化模型参数

from fast_pytorch_kmeans import KMeans
from collections import namedtuple
import pdb
Codebook = namedtuple('Codebook', ['centroids', 'labels'])
def k_means_quantize(fp32_tensor: torch.Tensor, bitwidth=4, codebook=None):"""quantize tensor using k-means clustering:param fp32_tensor::param bitwidth: [int] quantization bit width, default=4:param codebook: [Codebook] (the cluster centroids, the cluster label tensor):return:[Codebook = (centroids, labels)]centroids: [torch.(cuda.)FloatTensor] the cluster centroidslabels: [torch.(cuda.)LongTensor] cluster label tensor"""if codebook is None:# 首先计算聚类的中心点个数# get number of clusters based on the quantization precisionn_clusters = 2**bitwidth# print(n_clusters)# 用kmeans算法得到聚类的中心# use k-means to get the quantization centroidskmeans = KMeans(n_clusters=n_clusters, mode='euclidean', verbose=0)labels = kmeans.fit_predict(fp32_tensor.view(-1, 1)).to(torch.long)centroids = kmeans.centroids.to(torch.float).view(-1)codebook = Codebook(centroids, labels)# decode the codebook into k-means quantized tensor for inference# 解码codebook，得到k-means量化后的tensorquantized_tensor = codebook.centroids[codebook.labels]fp32_tensor.set_(quantized_tensor.view_as(fp32_tensor))return codebookdef plot_matrix(tensor, ax, title, cmap=ListedColormap(['white'])):ax.imshow(tensor.cpu().numpy(), vmin=-0.5, vmax=0.5, cmap=cmap)ax.set_title(title)ax.set_yticklabels([])ax.set_xticklabels([])for i in range(0,tensor.shape[0]):for j in range(0,tensor.shape[1]):text = ax.text(j, i, f'{tensor[i, j].item():.2f}',ha="center", va="center", color="k") if __name__ == "__main__":bitwidth = 2test_tensor = torch.tensor([[-0.3747,  0.0874,  0.3200, -0.4868,  0.4404],[-0.0402,  0.2322, -0.2024, -0.4986,  0.1814],[ 0.3102, -0.3942, -0.2030,  0.0883, -0.4741]])fig, axes = plt.subplots(1,2, figsize=(8, 12))ax_left, ax_right = axes.ravel()plot_matrix(test_tensor, ax_left, 'original tensor')num_unique_values_before_quantization = test_tensor.unique().numel()codebook_test = k_means_quantize(fp32_tensor=test_tensor, bitwidth=bitwidth)# pdb.set_trace()num_unique_values_after_quantization = test_tensor.unique().numel()print(f'    target bitwidth: {bitwidth} bits')print(f'        num unique values before k-means quantization: {num_unique_values_before_quantization}')print(f'        num unique values after  k-means quantization: {num_unique_values_after_quantization}')assert num_unique_values_after_quantization == min((1 << bitwidth), num_unique_values_before_quantization)plot_matrix(test_tensor, ax_right, f'{bitwidth}-bit k-means quantized tensor', \cmap='tab20c')fig.tight_layout()plt.show()

运行结果如下:

训练kmeans量化模型

accuracy_drop_threshold = 0.5
quantizers_before_finetune = copy.deepcopy(quantizers)
quantizers_after_finetune = quantizersfor bitwidth in [8, 4, 2]:quantizer = quantizers[bitwidth]print(f'k-means quantizing model into {bitwidth} bits')quantizer.apply(model, update_centroids=False)quantized_model_size = get_model_size(model, bitwidth)print(f"    {bitwidth}-bit k-means quantized model has size={quantized_model_size/MiB:.2f} MiB")quantized_model_accuracy = evaluate(model, test_loader)print(f"    {bitwidth}-bit k-means quantized model has accuracy={quantized_model_accuracy:.2f}% before quantization-aware training ")accuracy_drop = fp32_model_accuracy - quantized_model_accuracyif accuracy_drop > accuracy_drop_threshold:print(f"        Quantization-aware training due to accuracy drop={accuracy_drop:.2f}% is larger than threshold={accuracy_drop_threshold:.2f}%")num_finetune_epochs = 10optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_finetune_epochs)criterion = nn.CrossEntropyLoss()best_accuracy = 0epoch = num_finetune_epochswhile accuracy_drop > accuracy_drop_threshold and epoch > 0:train(model, train_loader, criterion, optimizer, scheduler,callbacks=[lambda: quantizer.apply(model, update_centroids=True)])model_accuracy = evaluate(model, test_loader)is_best = model_accuracy > best_accuracybest_accuracy = max(model_accuracy, best_accuracy)print(f'        Epoch {num_finetune_epochs-epoch} Accuracy {model_accuracy:.2f}% / Best Accuracy: {best_accuracy:.2f}%')accuracy_drop = fp32_model_accuracy - best_accuracyepoch -= 1else:print(f"        No need for quantization-aware training since accuracy drop={accuracy_drop:.2f}% is smaller than threshold={accuracy_drop_threshold:.2f}%")

训练过程截图:

参考资料

• 模型减肥秘籍:https://www.datawhale.cn/learn/content/68/963
• MIT教学课程:https://hanlab.mit.edu/courses/2024-fall-65940