motionbert模型量化后精度损失

针对 MotionBERT 这类时序 Transformer 模型，PTQ 遇到瓶颈时，QAT（量化感知训练）确实是最终解决方案。

社区已有成熟的 QAT 流程和配置指南，以下是针对你这种情况的QAT 落地路径：

核心参考资源

1. 必读精华帖

MotionBERT QAT 实施步骤

步骤 1：模型改造（关键）

MotionBERT 需要插入 QuantStub/DequantStub 并融合算子：

# quantize_prepare.py
import torch
import torch.nn as nn
from horizon_plugin_pytorch.quantization import QuantStub, DequantStub
from horizon_plugin_pytorch.nn.quantized import FloatFunctional

class MotionBERT_Quant(nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.model = original_model
        self.quant = QuantStub()
        self.dequant = DequantStub()
        
        # 替换 MatMul 为 FloatFunctional（Eager Mode 要求）
        self.matmul_qk = FloatFunctional()
        self.matmul_pv = FloatFunctional()
    
    def forward(self, x):
        x = self.quant(x)
        # 调用原始模型逻辑，但使用 self.matmul_qk.matmul(q, k) 替代 torch.matmul
        x = self.model(x)
        x = self.dequant(x)
        return x
    
    def fuse_model(self):
        # 融合 LayerNorm + Linear 等
        from horizon_plugin_pytorch.quantization import fuse_known_modules
        # 根据实际模型结构调整融合策略
        fuse_known_modules(self.model, [["norm1", "q_linear"], ["norm1", "k_linear"]])

步骤 2：配置 qconfig（针对 Transformer 优化）

# qat_config.py
from horizon_plugin_pytorch.quantization import get_default_qconfig
from horizon_plugin_pytorch.dtype import qint16
import torch

# 全局默认 int8
default_qconfig = get_default_qconfig(
    activation_fake_quant="fake_quant",
    weight_fake_quant="fake_quant",
    activation_observer="mse",  # Transformer 推荐 mse
    weight_observer="min_max",
)

# Attention 层专用 int16 配置
attention_qconfig = get_default_qconfig(
    activation_fake_quant="fake_quant",
    weight_fake_quant="fake_quant",
    activation_observer="mse",
    weight_observer="min_max",
    activation_qkwargs={"dtype": qint16},  # 激活用 int16
)

# 输出头 int32 配置
head_qconfig = get_default_qconfig(
    activation_fake_quant=None,  # 不量化输出
    weight_fake_quant="fake_quant",
    activation_observer=None,
    weight_observer="min_max",
)

# 构建 qconfig_dict
qconfig_dict = {
    "": default_qconfig,
    "module_name": {
        # 所有 Attention 层
        "model.blocks_st.0.attn_s.q_linear": attention_qconfig,
        "model.blocks_st.0.attn_s.k_linear": attention_qconfig,
        "model.blocks_st.0.attn_s.v_linear": attention_qconfig,
        "model.blocks_st.0.attn_t.q_linear": attention_qconfig,
        "model.blocks_st.0.attn_t.k_linear": attention_qconfig,
        "model.blocks_st.0.attn_t.v_linear": attention_qconfig,
        # ... 对所有 blocks_st 和 blocks_ts 重复配置
        
        # 输出头
        "model.head": head_qconfig,
    }
}

步骤 3：Calibration + QAT 训练

# train_qat.py
from horizon_plugin_pytorch.quantization import (
    prepare_qat, convert, set_fake_quantize, FakeQuantState
)
from horizon_plugin_pytorch.quantization.qconfig import (
    default_calib_8bit_fake_quant_qconfig,
    default_qat_8bit_fake_quant_qconfig
)
import torch

# 1. 加载浮点模型
float_model = MotionBERT_Quant(original_model)
float_model.fuse_model()
float_model.load_state_dict(torch.load("motionbert_float.pth"))
float_model.eval()

# 2. Calibration 阶段
float_model.qconfig = default_calib_8bit_fake_quant_qconfig
calib_model = prepare_qat(float_model, qconfig_dict)  # 使用自定义 qconfig_dict

set_fake_quantize(calib_model, FakeQuantState.CALIBRATION)
calib_model.eval()

# 用 300-500 条校准数据跑一遍
for data in calibration_loader:
    calib_model(data)

torch.save(calib_model.state_dict(), "calib_checkpoint.pth")

# 3. QAT 训练阶段
float_model.load_state_dict(torch.load("motionbert_float.pth"))
float_model.qconfig = default_qat_8bit_fake_quant_qconfig
qat_model = prepare_qat(float_model, qconfig_dict)
qat_model.load_state_dict(torch.load("calib_checkpoint.pth"))

set_fake_quantize(qat_model, FakeQuantState.QAT)
qat_model.train()

# 微调训练（关键超参）
optimizer = torch.optim.SGD(qat_model.parameters(), lr=0.001, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

for epoch in range(10):  # 通常 5-10 个 epoch 足够
    for data, target in train_loader:
        optimizer.zero_grad()
        output = qat_model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    
    scheduler.step()
    
    # 每个 epoch 验证一次
    qat_model.eval()
    set_fake_quantize(qat_model, FakeQuantState.VALIDATION)
    evaluate(qat_model, val_loader)
    qat_model.train()
    set_fake_quantize(qat_model, FakeQuantState.QAT)

torch.save(qat_model.state_dict(), "qat_checkpoint.pth")

# 4. 定点转换
qat_model.eval()
set_fake_quantize(qat_model, FakeQuantState.VALIDATION)
quantized_model = convert(qat_model)

# 验证定点精度
evaluate(quantized_model, val_loader)
torch.save(quantized_model.state_dict(), "quantized_checkpoint.pth")

# 5. 导出 ONNX
torch.onnx.export(
    quantized_model.cpu(),
    torch.randn(1, 120, 17, 3),
    "motionbert_qat.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=11
)

步骤 4：模型转换

# 使用 PTQ 相同的流程转换 QAT 导出的 ONNX
hb_mapper --config config_qat.yaml

注意：QAT 导出的 ONNX 通常不需要复杂的 node_info 配置，因为量化参数已嵌入模型。

QAT 超参推荐（针对时序模型）

常见问题

1. MatMul 不支持 int16？

QAT 中可以通过 qconfig 强制 Attention 层用 int16，工具链会在编译时自动处理。

2. 训练时 Nan？

• 降低 Learning Rate
• 使用 AMP 混合精度训练
• 检查校准数据是否有异常值

3. QAT 后精度仍不达标？

• 增加 QAT epoch 到 15-20
• 尝试 activation_observer="kl"
• 检查浮点模型是否已充分收敛

预期效果

下一步行动

1. 先阅读 QAT 快速上手，跑通示例代码
2. 参考 qconfig 使用指南，配置 MotionBERT 的 int16 策略
3. 准备原训练数据集，用于 QAT 微调
4. 从 Calibration 开始，验证精度提升后再进行 QAT 训练

如果需要更具体的代码适配帮助，可以提供：

• 你的 MotionBERT 模型定义代码
• 是否有原训练数据集
• 当前浮点模型的 MPJPE 误差

QAT 虽然实施复杂度高于 PTQ，但对于 MotionBERT 这类时序 Transformer 模型，这是获得生产级精度的必经之路。