pytorch模型int8量化 MNIST

### PyTorch MNIST 模型 INT8 量化的实现

对于希望减少计算资源消耗并提高推理速度的应用场景而言，INT8量化是一种有效的方法。通过降低权重和激活值的精度到8位整数表示，可以在保持较高准确性的同时显著提升性能。

在PyTorch中执行INT8量化涉及几个重要步骤：

#### 准备环境与加载预训练模型
为了确保后续操作顺利进行，先安装必要的库，并导入所需的模块[^1]。

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from tqdm import tqdm

接着定义转换函数以及数据集加载器来获取MNIST测试集用于校准过程：

transform = transforms.Compose([
    transforms.ToTensor(),
])

test_dataset = datasets.MNIST(root='./data', train=False,
                              download=True, transform=transform)

calibration_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

#### 定义辅助类来进行统计收集
创建一个简单的观察者类以记录每一层的最大最小值范围，这对于后续确定合适的缩放因子至关重要：

class MinMaxObserver(torch.nn.Module):
    def __init__(self):
        super(MinMaxObserver, self).__init__()
        self.min_val = float('inf')
        self.max_val = -float('inf')

    def forward(self, x):
        min_x = torch.min(x).item()
        max_x = torch.max(x).item()

        if min_x < self.min_val:
            self.min_val = min_x
        if max_x > self.max_val:
            self.max_val = max_x
            
        return x

#### 应用量化感知训练 (QAT) 或静态量化方法
这里展示的是基于静态仿真的方式，在此之前需确保已有一个经过充分训练好的浮点版本模型实例 `model` 可供使用。应用量化配置前记得设置为评估模式：

model.eval()  # Switch the model into evaluation mode.
quantized_model = torch.quantization.convert(model.to('cpu'), inplace=False)

如果采用动态量化，则只需指定哪些类型的层需要被处理；而对于静态量化来说，除了上述之外还需要额外提供代表性的输入样本以便于调整比例尺参数：

# Static Quantization Preparation
fused_model = torch.quantization.fuse_modules(
    copy.deepcopy(model), [['conv1', 'relu1'], ['conv2', 'relu2']])
qconfig_dict = {"": torch.quantization.get_default_qat_qconfig('fbgemm')}
prepared_quantized_model = torch.quantization.prepare_qat(fused_model, qconfig_spec=qconfig_dict)

for images, _ in calibration_loader:
    prepared_quantized_model(images)

final_quantized_model = torch.quantization.convert(prepared_quantized_model.cpu().eval(), inplace=False)

完成以上流程之后便得到了适用于部署阶段使用的低比特宽度网络结构——即完成了从FP32至INT8的数据类型转变工作[^2]。