基于示例详细讲解模型PTQ量化的步骤（含代码）

通过上述详细步骤，我们对一个训练好的CNN模型进行了PTQ。详细步骤包括从模型中收集统计信息、选择和计算量化参数、应用量化到权重和激活、进行模型校准以及最终的模型验证和评估。每个步骤涉及具体的计算和调整，以确保量化过程中的模型性能尽可能接近原始浮点模型。

LQS2020

1921人浏览 · 2024-09-02 10:52:12

LQS2020 · 2024-09-02 10:52:12 发布

详细探讨模型PTQ量化每个步骤，涉及更多的技术细节和实际计算方法，以便更好地理解PTQ（Post-Training Quantization，训练后量化）的全过程。

1. 模型训练

我们假设已经训练了一个卷积神经网络（CNN），例如VGG-16。训练完成后，我们得到了一个以32位浮点数表示的模型权重和激活值。

2. 收集统计信息

在量化之前，我们需要从模型中收集统计信息，以帮助确定量化的参数。

收集权重和激活的统计信息

1. 权重统计

对于每个卷积层和全连接层：

最大值和最小值：通过遍历模型的每个权重矩阵，计算权重的最大值和最小值。

import numpy as np

def get_weight_stats(weights):
    max_val = np.max(weights)
    min_val = np.min(weights)
    return max_val, min_val

示例：

weights_conv1 = model.conv1.weight.data.numpy()  # 获取卷积层1的权重
max_weight_conv1, min_weight_conv1 = get_weight_stats(weights_conv1)

2. 激活统计

激活值的统计信息通常在校准过程中收集：

最大值和最小值：将校准数据集通过模型，记录每个层的激活值的统计信息。

def get_activation_stats(model, dataloader):
    activations = []
    for inputs, _ in dataloader:
        outputs = model(inputs)
        activations.append(outputs.detach().numpy())
    max_activation = np.max(activations)
    min_activation = np.min(activations)
    return max_activation, min_activation

示例：

max_activation, min_activation = get_activation_stats(model, calibration_dataloader)

3. 选择量化方案

根据收集到的统计数据，选择量化方案并计算量化参数。

选择量化位宽

通常选择8位整数（INT8），即量化到[-128, 127]范围内。

计算量化参数

1. 对称量化

对于权重：

计算缩放因子：

def calculate_scale(min_val, max_val, int_min, int_max):
    scale = (max_val - min_val) / (int_max - int_min)
    return scale

量化公式：

def quantize_weight(weights, scale, int_min, int_max):
    quantized_weights = np.clip(np.round(weights / scale), int_min, int_max)
    return quantized_weights

示例：

scale_weight = calculate_scale(min_weight_conv1, max_weight_conv1, -128, 127)
quantized_weights_conv1 = quantize_weight(weights_conv1, scale_weight, -128, 127)

2. 非对称量化

对于激活值：

计算缩放因子和零点：

def calculate_activation_params(min_val, max_val, int_min, int_max):
    scale = (max_val - min_val) / (int_max - int_min)
    zero_point = int_min - np.round(min_val / scale)
    return scale, zero_point

量化公式：

def quantize_activation(activations, scale, zero_point, int_min, int_max):
    quantized_activations = np.clip(np.round(activations / scale) + zero_point, int_min, int_max)
    return quantized_activations

示例：

scale_activation, zero_point = calculate_activation_params(min_activation, max_activation, 0, 255)
quantized_activations = quantize_activation(activation_data, scale_activation, zero_point, 0, 255)

4. 量化权重

权重量化步骤：

计算缩放因子：

scale = calculate_scale(min_weight, max_weight, -128, 127)

应用量化公式：

quantized_weights = quantize_weight(weights, scale, -128, 127)

存储量化参数：

保存量化的缩放因子和偏移量，这在推理阶段用于反量化。
```
np.save('quantized_weights.npy', quantized_weights)
np.save('weight_scale.npy', scale)
```

5. 量化激活

激活量化步骤：

计算激活的缩放因子和零点：

scale, zero_point = calculate_activation_params(min_activation, max_activation, 0, 255)

应用量化公式：

quantized_activations = quantize_activation(activations, scale, zero_point, 0, 255)

存储量化参数：

保存激活的量化参数，用于反量化。

np.save('activation_scale.npy', scale)
np.save('activation_zero_point.npy', zero_point)

6. 模型校准

微调（Fine-Tuning）：

步骤：

将量化后的模型加载并用量化参数初始化。
用量化后的模型和校准数据集进行轻微的训练，以优化量化效果。

from torch.optim import Adam

# 定义优化器和损失函数
optimizer = Adam(model.parameters(), lr=1e-5)
criterion = torch.nn.CrossEntropyLoss()

# 轻微训练
model.train()
for epoch in range(1):
    for inputs, targets in calibration_dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

7. 验证和评估

步骤：

测试量化模型：

使用测试数据集对量化后的模型进行评估，比较其与原始浮点模型的性能。

model.eval()
correct = 0
total = 0
with torch.no_grad():
    for inputs, targets in test_dataloader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += targets.size(0)
        correct += (predicted == targets).sum().item()

accuracy = correct / total
print(f'Test Accuracy: {accuracy * 100:.2f}%')