量化Quantization初步之--带量化(QAT)的XOR异或pyTorch版250501

量化(Quantization)这词儿听着玄，经常和量化交易Quantitative Trading (量化交易)混淆。

其实机器学习(深度学习)领域的量化Quantization是和节约内存、提高运算效率相关的概念（因大模型的普及，这个量化问题尤为迫切）。

揭秘机器学习“量化”：不止省钱，更让AI高效跑起来！

“量化（Quantization）”这个词，在机器学习领域，常常让人联想到复杂的数学或是与金融交易相关的“量化交易”，从而感到困惑。但实际上，它与我们日常生活中的数字转换概念更为接近，而在AI世界里，它扮演的角色是节约内存、提高运算效率的“幕后英雄”(现在已经显露到“幕布之前”），尤其在大模型时代，其重要性日益凸显。

那么，机器学习中的“量化”究竟是啥？咱为啥用它？

什么是机器学习中的“量化”（Quantization）？

简单讲，机器学习中的“量化”就是将模型中原本采用高精度浮点数（如32位浮点数，即FP32）表示的权重（weights）和激活值（activations），转换成低精度表示（如8位整数，即INT8）的过程。

你可以把它想象成“数字的压缩”。在计算机中，浮点数就像是拥有无限小数位的精确数字，而整数则像只有整数部分的数字。从高精度浮点数到低精度整数的转换，必然会损失一些信息，但与此同时，它也带来了显著的优势：

1. 内存占用大幅减少： 8位整数比32位浮点数少占用4倍的内存空间。这意味着更大的模型可以被部署到内存有限的设备上（如手机、IoT设备），或者在相同内存下可以运行更大的模型。
2. 计算速度显著提升： 整数运算通常比浮点数运算更快、功耗更低。这使得模型在推理（Inference）阶段能以更高的效率运行，减少延迟。

为何需要“量化”？

随着深度学习模型变得越来越大，越来越复杂，它们对计算资源的需求也呈爆炸式增长。一个动辄几十亿甚至上百亿参数的大模型，如果全部使用FP32存储和计算，将对硬件资源提出极高的要求。

• 部署到边缘设备： 手机、自动驾驶汽车、智能音箱等边缘设备通常算力有限，内存紧张。量化是让大模型“瘦身”后成功“登陆”这些设备的必经之路。
• 降低运行成本： 在云端部署大模型时，更低的内存占用和更快的计算速度意味着更低的服务器成本和能耗。
• 提升用户体验： 实时响应的AI应用，如语音助手、图像识别等，对推理速度有极高要求。量化可以有效缩短响应时间。

量化策略：后训练量化 vs. 量化感知训练（QAT）

量化并非只有一种方式。根据量化发生的时间点，主要可以分为两大类：

1. 后训练量化（Post-Training Quantization, PTQ）： 顾名思义，PTQ 是在模型训练完成之后，对已经训练好的FP32模型进行量化。它操作简单，不需要重新训练，是实现量化的最快途径。然而，由于量化过程中会损失精度，PTQ 可能会导致模型性能（如准确率）的下降。对于对精度要求不那么苛刻的应用，PTQ 是一个不错的选择。

2. 量化感知训练*（这是本文重点推介的：Quantization Aware Training, QAT）： 这正是我们今天着重讲解的明星策略！QAT 的核心思想是——在模型训练过程中，就“感知”到未来的量化操作。

   在QAT中，量化误差被集成到模型的训练循环中。这意味着，模型在训练时就“知道”它最终会被量化成低精度，并会努力学习如何在这种低精度下保持最优性能。

   具体来说，QAT通常通过在模型中插入“伪量化”（Fake Quantization）节点来实现。这些节点在训练过程中模拟量化和反量化操作，使得模型在FP32环境下进行前向传播和反向传播时，能够学习到量化对模型参数和激活值的影响。当训练完成后，这些伪量化节点会被真正的量化操作所取代，从而得到一个高性能的量化模型。

为什么QAT是量化策略的“王牌”？

相较于PTQ……QAT 的优势显而易见：

• 精度损失最小： 这是QAT最大的亮点。通过在训练过程中模拟量化，模型能够自我调整以适应量化带来的精度损失，从而在量化后依然保持接近FP32模型的性能。
• 适用于更苛刻的场景： 对于那些对模型精度要求极高，不能容忍明显性能下降的应用（如自动驾驶、医疗影像分析），QAT几乎是唯一的选择。
• 更好的泛化能力： 在训练阶段就考虑量化，使得模型在量化后对各种输入数据具有更好的鲁棒性。

PyTorch中的QAT实践

在PyTorch中实现QAT，通常需要以下几个关键步骤：

1. 准备量化配置： 定义量化类型（如INT8）、量化方法（如对称量化、非对称量化）以及需要量化的模块。
2. 模型转换： 使用PyTorch提供的 torch.quantization 模块，将普通的FP32模型转换为QAT模型。这个过程会在模型中插入伪量化模块。
3. 重新训练/微调： 在新的数据集上对转换后的模型进行短时间的微调（Fine-tuning），或者在原有训练基础上继续训练。这个阶段，模型会学习如何适应伪量化带来的精度损失。
4. 模型融合（可选但推荐）： 将一些连续的层（如Conv-BN-ReLU）融合为一个操作，可以进一步提高量化后的推理效率。
5. 模型量化和保存： 训练完成后，将微调好的QAT模型转换为真正的量化模型，并保存。

总结

量化（Quantization）是深度学习模型优化不可或缺的一环，它通过降低模型精度来换取内存和计算效率的大幅提升。而量化感知训练（QAT）作为一种高级量化策略，通过在训练阶段就考虑量化对模型的影响，极大地减小了量化带来的精度损失，使得在各种设备上部署高性能AI模型成为可能。

随着大模型和边缘AI的普及，掌握量化尤其是QAT的原理和实践，将成为每一位AI工程师和研究人员的必备技能。让我们一起，让AI跑得更快、更高效！

import torch
import torch.nn as nn
import torch.optim as optim
import torch.quantization
import numpy as np
import os# ===== 1. XOR数据集 =====
X = torch.tensor([[0., 0.],[0., 1.],[1., 0.],[1., 1.]
], dtype=torch.float32)
y = torch.tensor([[0.],[1.],[1.],[0.]
], dtype=torch.float32)# ===== 2. 神经网络模型 (标准FP32) =====
class XORNet(nn.Module):def __init__(self):super(XORNet, self).__init__()self.fc1 = nn.Linear(2, 3)self.relu = nn.ReLU()self.fc2 = nn.Linear(3, 1)# QAT阶段不用Sigmoid，直接用BCEWithLogitsLoss# self.sigmoid = nn.Sigmoid()def forward(self, x):x = self.fc1(x)x = self.relu(x)x = self.fc2(x)# return self.sigmoid(x)return x# ===== 3. 初始化模型/优化器 =====
model = XORNet()
# He初始化，适合ReLU
for m in model.modules():if isinstance(m, nn.Linear):nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='relu')if m.bias is not None:nn.init.constant_(m.bias, 0)criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.05)# ===== 4. 训练（标准模型） =====
print("--- 开始标准模型训练 ---")
epochs = 900#1000#500    #1500
for epoch in range(epochs):outputs = model(X)loss = criterion(outputs, y)optimizer.zero_grad()loss.backward()optimizer.step()if (epoch + 1) % 300 == 0:print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')with torch.no_grad():probs = torch.sigmoid(model(X))predictions = (probs > 0.5).float()accuracy = (predictions == y).sum().item() / y.numel()print(f"\n标准模型训练后精度: {accuracy*100:.2f}%")print(f"标准模型预测结果:\n{predictions}")# ===== 5. 构建QAT模型 =====
class XORNetQAT(nn.Module):def __init__(self):super(XORNetQAT, self).__init__()# 量化Stubself.quant = torch.quantization.QuantStub()self.fc1 = nn.Linear(2, 3)self.relu = nn.ReLU()self.fc2 = nn.Linear(3, 1)self.dequant = torch.quantization.DeQuantStub()def forward(self, x):x = self.quant(x)x = self.fc1(x)x = self.relu(x)x = self.fc2(x)x = self.dequant(x)return xdef fuse_model(self):torch.quantization.fuse_modules(self, [['fc1', 'relu']], inplace=True)# ===== 6. QAT前权重迁移、模型融合 =====
model_qat = XORNetQAT()
# 迁移参数
model_qat.load_state_dict(model.state_dict())
# 融合（此步必须！）
model_qat.fuse_model()# 配置QAT
model_qat.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')  # CPU
torch.quantization.prepare_qat(model_qat, inplace=True)optimizer_qat = optim.Adam(model_qat.parameters(), lr=0.01)# ===== 7. QAT训练 =====
qat_epochs = 700
print("\n--- QAT训练 ---")
for epoch in range(qat_epochs):model_qat.train()outputs_qat = model_qat(X)loss_qat = criterion(outputs_qat, y)optimizer_qat.zero_grad()loss_qat.backward()optimizer_qat.step()if (epoch + 1) % 150 == 0:print(f'QAT Epoch [{epoch+1}/{qat_epochs}], Loss: {loss_qat.item():.4f}')# ===== 8. 转换到量化模型/评估精度 =====
print("\n--- 转换为量化模型 ---")
model_qat.eval()
model_quantized = torch.quantization.convert(model_qat.eval(), inplace=False)with torch.no_grad():probs_quantized = torch.sigmoid(model_quantized(X))predictions_quantized = (probs_quantized > 0.5).float()accuracy_quantized = (predictions_quantized == y).sum().item() / y.numel()print(f"量化模型精度: {accuracy_quantized*100:.2f}%")print(f"量化模型预测:\n{predictions_quantized}")# ===== 9. 模型大小对比 =====
torch.save(model.state_dict(), 'xor_fp32.pth')
torch.save(model_quantized.state_dict(), 'xor_int8.pth')
fp32_size = os.path.getsize('xor_fp32.pth') / (1024 * 1024)
int8_size = os.path.getsize('xor_int8.pth') / (1024 * 1024)
print(f"\nFP32模型大小: {fp32_size:.6f} MB")
print(f"INT8模型大小: {int8_size:.6f} MB")
print(f"模型缩减比例: {fp32_size/int8_size:.2f} 倍")

运行结果：

====================== RESTART: F:/qatXorPytorch250501.py 
--- 开始标准模型训练 ---
Epoch [300/900], Loss: 0.0033
Epoch [600/900], Loss: 0.0011
Epoch [900/900], Loss: 0.0005

标准模型训练后精度: 100.00%
标准模型预测结果:
tensor([[0.],
        [1.],
        [1.],
        [0.]])

--- QAT训练 ---
QAT Epoch [150/700], Loss: 0.0005
QAT Epoch [300/700], Loss: 0.0004
QAT Epoch [450/700], Loss: 0.0004
QAT Epoch [600/700], Loss: 0.0003

--- 转换为量化模型 ---
量化模型精度: 100.00%
量化模型预测:
tensor([[0.],
        [1.],
        [1.],
        [0.]])

FP32模型大小: 0.001976 MB
INT8模型大小: 0.004759 MB
模型缩减比例: 0.42 倍

最后：
其实，这次量化
量化后的模型是量化前的 4.2倍（咦？不是说好了压缩吗?咋变大了?)


魔鬼藏在细节：
咱们看看 量化 之前的 (基线的）模型 参数+权重等等：

===== Model Architecture =====
XORNet(
  (fc1): Linear(in_features=2, out_features=3, bias=True)
  (relu): ReLU()
  (fc2): Linear(in_features=3, out_features=1, bias=True)
)

===== Layer Parameters =====
[fc1.weight] shape: (3, 2)
[[ 1.723932,  1.551827],
 [ 2.106917,  1.681809],
 [-0.299378, -0.444912]]

[fc1.bias] shape: (3,)
[-1.725313, -2.509506,  0.      ]

[fc2.weight] shape: (1, 3)
[[-2.492318, -3.94821 ,  0.911841]]

[fc2.bias] shape: (1,)
[0.692789]

===== Extra Info (Hyperparameters) =====
Optimizer: Adam
Learning Rate: 0.05
Epochs: 900
Loss: BCEWithLogitsLoss
Activation: ReLU

量化之后的模型参数等：
超参数部分：
 

===== Model Architecture =====
XORNetQAT(
  (quant): Quantize(scale=tensor([0.0079]), zero_point=tensor([0]), dtype=torch.quint8)
  (fc1): QuantizedLinearReLU(in_features=2, out_features=3, scale=0.0678500160574913, zero_point=0, qscheme=torch.per_channel_affine)
  (relu): Identity()
  (fc2): QuantizedLinear(in_features=3, out_features=1, scale=0.1376650333404541, zero_point=65, qscheme=torch.per_channel_affine)
  (dequant): DeQuantize()
)

===== Layer Parameters =====
===== Extra Info (Hyperparameters) =====
Optimizer: Adam
Learning Rate: 0.01
QAT Epochs: 700
Loss: BCEWithLogitsLoss
QConfig: QConfig(activation=functools.partial(<class 'torch.ao.quantization.fake_quantize.FusedMovingAvgObsFakeQuantize'>, observer=<class 'torch.ao.quantization.observer.MovingAverageMinMaxObserver'>, quant_min=0, quant_max=255, reduce_range=True){'factory_kwargs': <function _add_module_to_qconfig_obs_ctr.<locals>.get_factory_kwargs_based_on_module_device at 0x00000164FDB16160>}, weight=functools.partial(<class 'torch.ao.quantization.fake_quantize.FusedMovingAvgObsFakeQuantize'>, observer=<class 'torch.ao.quantization.observer.MovingAveragePerChannelMinMaxObserver'>, quant_min=-128, quant_max=127, dtype=torch.qint8, qscheme=torch.per_channel_symmetric){'factory_kwargs': <function _add_module_to_qconfig_obs_ctr.<locals>.get_factory_kwargs_based_on_module_device at 0x00000164FDB16160>})

看到吗？
量化后， 参数 变多了哈哈!

那量化的意义到底在哪里呢？？

在下面的 参数（非超参）的 权重的部分：

===== 量化 (Quantization) 后参数 =====

[fc1] (QuantizedLinearReLU)
  [weight] shape: torch.Size([3, 2]), dtype: torch.qint8
  weight (quantized):
[[ 3.97  3.98]
 [ 3.95  3.96]
 [-0.56 -1.05]]
  weight (raw int):
[[127 127]
 [127 127]
 [-18 -33]]
  scale: 0.06785
  zero_point: 0
  [bias] shape: torch.Size([3]), dtype: torch.float
  bias:
[-3.049073e-04 -3.960315e+00  0.000000e+00]

[fc2] (QuantizedLinear)
  [weight] shape: torch.Size([1, 3]), dtype: torch.qint8
  weight (quantized):
[[ 3.77 -7.79  0.41]]
  weight (raw int):
[[ 27 -57   3]]
  scale: 0.13766
  zero_point: 65
  [bias] shape: torch.Size([1]), dtype: torch.float
  bias:
[-7.147094]

再看看量化前：

===== Layer Parameters =====
[fc1.weight] shape: (3, 2)
[[ 1.723932,  1.551827],
 [ 2.106917,  1.681809],
 [-0.299378, -0.444912]]

[fc1.bias] shape: (3,)
[-1.725313, -2.509506,  0.      ]

[fc2.weight] shape: (1, 3)
[[-2.492318, -3.94821 ,  0.911841]]

[fc2.bias] shape: (1,)
[0.692789]

最后看看 量化 后：

===== 量化 (Quantization) 后参数 =====

[fc1] (QuantizedLinearReLU)
  [weight] shape: torch.Size([3, 2]), dtype: torch.qint8
  weight (quantized):
[[ 3.97  3.98]
 [ 3.95  3.96]
 [-0.56 -1.05]]
  weight (raw int):
[[127 127]
 [127 127]
 [-18 -33]]
  scale: 0.06785
  zero_point: 0
  [bias] shape: torch.Size([3]), dtype: torch.float
  bias:

看出区别了吗？

So：

量化操作 只 适用于 大、中型号的模型……道理就在此：
量化前 的 Weights 的权重 全部都是： Float32浮点型 ……很占内存的！
量化后 是所谓INT8（即一个字节、8bits）……至少（在权重部分）节省了 3/4 的内存！

So： 大模型 必须 要 量化 才 节省内存。
当然前提 是 你的 GPU 硬件 要 支持 INT8(8bits)的 运算哦……（这是后话，下次再聊）。

(over)完！