QAT 实战:从 PyTorch 模型到 J6E 实板精度与延时验证

本文从 CIFAR-10 浮点模型开始,依次完成校准、QAT 微调、BC 导出、HBM 编译,最后通过 UCP 在 J6E 开发板上验证精度和推理延时。整个流程基于 OpenExplorer 3.8.1,目标架构为 nash-e。

1. 实验环境与结果

实验使用的主要环境如下:

项目 配置
开发板 Horizon Robotics J6E MatrixB
BPU 架构 nash-e
BPU 核心 Core 0
开发环境 OpenExplorer 3.8.1 GPU Docker
PyTorch 2.8.0+cu128
horizon-plugin-pytorch 3.2.4
HBDK4 Compiler 4.9.7
部署方式 x86 端通过 hbm_infer 调用板端 BPU

各阶段精度如下:

阶段 Top-1 精度变化
Float PyTorch 84.40% 基线
Calibration 84.26% -0.14pp
QAT PyTorch 86.03% 相对 Float +1.63pp
QAT BC 85.91% 相对 QAT -0.12pp
Quantized BC 85.89% -0.02pp
J6E HBM 85.89% 相对 Quantized BC 0.00pp

延时测试使用固定输入,预热 100 帧后正式测量 1,000 帧:

指标 Mean P50 P95 P99
x86 UCP 端到端 2.133 ms 2.135 ms 2.320 ms 2.503 ms
gRPC 发送到接收 1.765 ms 1.762 ms 1.928 ms 2.130 ms
板端总处理 0.633 ms 0.616 ms 0.762 ms 0.947 ms
BPU 纯推理 0.445 ms 0.430 ms 0.491 ms 0.740 ms

2. 整体流程

本次实践的主线如下:

arduino 复制代码
CIFAR-10 数据
    ↓
训练浮点 MobileNetV2
    ↓
prepare 量化计算图
    ↓
HistogramObserver 校准
    ↓
QAT 微调
    ↓
导出 QAT BC
    ↓
转换为 Quantized BC
    ↓
remove_io_op,改为整数输入输出
    ↓
编译 nash-e HBM
    ↓
通过 hbm_infer 部署到 J6E
    ↓
实板精度验证与延时测试

3. 数据预处理必须保持一致

本次使用 CIFAR-10 数据集:

  • 训练集 50,000 张
  • 测试集 10,000 张
  • 图像尺寸为 32×32
  • 共 10 个类别

模型使用的输入预处理为:

scss 复制代码
rgb_tensor = torch.from_numpy(rgb).permute(0, 3, 1, 2).contiguous()
centered_yuv = rgb2centered_yuv(rgb_tensor).to(torch.float32)
centered_yuv.div_(128.0)

对应的数据变化是:

go 复制代码
RGB uint8 NHWC
→ RGB uint8 NCHW
→ centered YUV444
→ float32 / 128

训练、BC 验证和实板验证必须执行完全相同的预处理。预处理不一致时,模型通常仍能正常运行,但输出精度会明显异常。

这类问题很容易被误判为量化误差或板端计算误差,因此应在部署前固定预处理流程,并在每个验证阶段复用同一份实现。

4. 模型的 QAT 适配

实验模型选择 MobileNetV2。它规模适中,同时包含普通卷积、深度可分离卷积和残差连接,适合作为 QAT 入门模型。

4.1 设置目标架构

构建模型、执行 prepare、导出和编译之前,都需要设置正确的 march:

java 复制代码
from horizon_plugin_pytorch import set_march

set_march("nash-e")

march 必须与目标板卡的 BPU 架构一致。

4.2 添加量化边界

模型输入和输出位置分别加入 QuantStub 和 DeQuantStub:

ruby 复制代码
self.quant = QuantStub()
self.dequant = torch.ao.quantization.DeQuantStub()

def forward(self, x):
    x = self.quant(x)
    x = self._forward_impl(x)
    return self.dequant(x)

QuantStub 标记浮点输入进入量化计算图的位置,DeQuantStub 标记输出恢复为浮点的位置。

后续执行 remove_io_op 后,这两个边界算子会被移除,最终 HBM 直接接收 INT8 输入并返回 INT32 输出。

4.3 处理 ReLU6 兼容性

torchvision 的 MobileNetV2 默认包含 ReLU6(inplace=True)。当前使用的 Horizon QAT 模块不支持这种形式,prepare 时会报错:

ini 复制代码
RuntimeError: Do not support ReLU6(inplace=True).

可以在 prepare 前将其改为非原地操作:

arduino 复制代码
for module in self.modules():
    if isinstance(module, nn.ReLU6):
        module.inplace = False

ReLU6 没有可训练参数,因此这项修改不会改变 checkpoint 的 state_dict 结构。

5. QConfig、Template 和 prepare

在 Horizon QAT 流程中,可以这样理解三个概念:

ini 复制代码
QConfig       = 具体怎样量化
Template      = 量化配置应用到哪里
QConfigSetter = 汇总规则并设置到模型

QConfig 决定 Observer 类型、数据类型、量化粒度和对称方式。例如:

ini 复制代码
qconfig = QConfig(
    weight=FakeQuantize.with_args(
        observer=MinMaxObserver,
        dtype=qint8,
        qscheme=torch.per_channel_symmetric,
        ch_axis=0,
    ),
    output=FakeQuantize.with_args(
        observer=MinMaxObserver,
        dtype=qint8,
        qscheme=torch.per_tensor_symmetric,
        ch_axis=-1,
    ),
)

Template 则负责按模块名或输入输出位置匹配量化配置。本次实验使用全局双 INT8 配置,没有加入 INT16 或 FP16 混合精度策略。

prepare 会分析浮点计算图,并完成以下工作:

  1. 根据 Template 计算各位置的最终 QConfig。
  2. 将浮点模块替换为 QAT 兼容模块。
  3. 插入 Observer 和 FakeQuantize 节点。
  4. 检查量化边界、模块执行次数和异常配置。
  5. 生成能够在校准、QAT 和验证状态之间切换的模型。

prepare 后得到的仍然是 PyTorch 模型,但前向计算已经可以模拟定点量化误差。

常见状态包括:

  • CALIBRATION:收集激活统计信息。
  • QAT:训练过程中启用 FakeQuant。
  • VALIDATION:冻结统计信息并验证当前量化精度。
  • FLOAT:关闭 FakeQuant,用于浮点行为检查。

执行后应检查 QConfig 变更日志和模型检查报告,确认没有异常配置或非预期量化边界。

6. 浮点训练、校准与 QAT

6.1 浮点模型

浮点模型使用 Adam 训练 100 个 epoch:

复制代码
optimizer       Adam
learning_rate   0.001
weight_decay    0.001
batch_size      256
epochs          100

最终 Top-1 为:

erlang 复制代码
84.40%

QAT 不应替代正常的浮点训练。更稳妥的做法是先得到精度可用的浮点模型,再使用更小的学习率进行量化微调。

6.2 校准

校准阶段使用 HistogramObserver,共执行 100 个 step,覆盖 25,600 张图像。

结果如下:

css 复制代码
Float Top-1         84.40%
Calibration Top-1   84.26%
精度变化            -0.14pp

校准阶段不通过梯度更新模型权重,主要作用是收集激活分布并确定量化范围。0.14 个百分点的下降处于可接受范围内,说明当前 INT8 配置可以继续进入 QAT。

6.3 QAT 微调

QAT 阶段切换到 MinMaxObserver,使用较小学习率训练 20 个 epoch:

复制代码
全局量化配置       权重 INT8 + 激活 INT8
optimizer          SGD
learning_rate      0.0001
weight_decay       0.0002
epochs             20

最终结果:

css 复制代码
QAT Top-1    86.03%
相对 Float   +1.63pp

QAT 精度高于浮点基线不代表量化天然能够提升模型精度。本次浮点模型在第 100 个 epoch 仍然取得最佳结果,说明训练尚未完全进入平台期。QAT 又增加了 20 个 epoch,同时改变了优化器和训练噪声,因此精度提升更可能来自额外微调和 FakeQuant 带来的正则化效果。

7. 从 QAT 模型导出 BC

QAT 模型导出后得到:

复制代码
mobilenetv2_cifar10_qat.bc

导出过程中曾遇到一个值得记录的问题:文件保存成功后,进程在退出阶段出现原生库段错误。

java 复制代码
Segmentation fault (core dumped)

仅凭这条日志无法判断 BC 是否损坏。随后我在一个全新的进程中重新加载 BC,并检查函数数量和输入输出签名:

lua 复制代码
HBDK4 load       成功
function_count   1
进程退出码       0

函数签名为:

go 复制代码
func @forward(tensor<1x3x32x32xf32> _input_0)
  -> tensor<1x10xf32> _output_0

因此该 BC 可以继续使用。

这里的经验是:产物已经生成,不等于产物一定有效;进程异常退出,也不等于产物一定损坏。更可靠的方法是在独立进程中重新加载产物并执行结构检查。

8. Quantized BC 与整数 IO

QAT BC 经过 convert(..., march="nash-e") 后得到 Quantized BC。转换前后精度如下:

css 复制代码
QAT BC Top-1          85.91%
Quantized BC Top-1    85.89%
精度变化              -0.02pp

转换带来的精度变化很小。

接着使用 remove_io_op 删除输入 Quantize 和输出 DeQuantize,使模型直接使用整数输入输出。处理后的函数签名为:

go 复制代码
func @forward(tensor<1x3x32x32xsi8> _input_0)
  -> tensor<1x10xsi32> _output_0

算子统计中没有发现 hbtl CPU 算子,主体计算均落在 BPU 上。

执行 remove_io_op 后,输入量化和输出反量化不再由模型内部完成,而是交给应用侧处理。

9. 编译 HBM

使用 HBDK4 编译面向 nash-e 的 HBM:

lua 复制代码
march          nash-e
optimization   2
debug          false

编译结果:

python 复制代码
compile_exit_code   0
compile_elapsed     17.69s
hbm_size            1,834,704 bytes

生成的 HBM 为:

复制代码
mobilenetv2_cifar10_nash_e.hbm

编译完成后,建议保存 HBM 的文件大小和 SHA256。这样可以确认实板测试使用的模型是否与编译产物一致。

10. 通过 hbm_infer 部署到 J6E

本次使用 HbmRpcSession 从 x86 端自动部署 HBM,并通过 gRPC 调用板端 BPU:

ini 复制代码
from hbm_infer.hbm_rpc_session import HbmRpcSession

session = HbmRpcSession(
    host="192.168.0.65",
    local_hbm_path="mobilenetv2_cifar10_nash_e.hbm",
    username="root",
    password=None,
    remote_root="/tmp/hbm_infer/",
    with_profile=True,
    compress_option="NONE",
    core_id=0,
)

try:
    outputs = session(inputs, model_name="forward")
finally:
    session.close_server()

hbm_infer 会自动完成:

  1. 通过 SSH 连接开发板。
  2. 部署 HBM、板端服务和运行库。
  3. 启动推理服务。
  4. 通过 gRPC 发送输入并接收结果。
  5. 关闭服务并清理临时文件。

HBM 不需要提前手工复制到开发板。

临时目录不要放在 NFS 上

最初我将 remote_root 放在 NFS 目录中。推理能够正常完成,但 close_server() 清理文件时出现:

php 复制代码
rm: cannot remove ...: Directory not empty

将临时目录改到板端本地文件系统后,问题消失:

bash 复制代码
/tmp/hbm_infer/

对于频繁创建和删除临时文件的推理服务,板端本地 /tmp 比 NFS 目录更合适。

无论推理是否成功,都应在 finally 中调用 close_server(),避免板端残留服务进程和临时文件。

11. 整数输入输出的量化处理

从 HBM 中读取到的输入信息如下:

css 复制代码
输入形状     [1, 3, 32, 32]
输入类型     DATA_TYPE_S8
输入 scale   0.006862768437713385
输入 zero    0

应用侧需要按照模型的量化参数将浮点输入转换为 INT8:

ini 复制代码
q = clamp(round(x / scale + zero_point), -128, 127)

输出信息如下:

ini 复制代码
输出形状     [1, 10]
输出类型     DATA_TYPE_S32
量化方式     沿 axis=1 逐通道量化

10 个类别分别使用独立的 scale,因此输出需要逐通道反量化:

css 复制代码
logit[c] = (output_int32[c] - zero_point[c]) × scale[c]

这里有一个容易遗漏的问题:不能直接对原始 INT32 输出执行 argmax。因为每个类别使用的 scale 不同,整数值之间不能直接比较。必须先逐通道反量化,再计算分类结果。

12. 实板精度验证

正式评测前,可以先用少量样本检查整个链路:

arduino 复制代码
python3 ucp_eval_cifar10.py --max-samples 100

小样本结果只用于确认预处理、输入量化、推理和输出解析是否正常。

在 CIFAR-10 的 10,000 张测试图像上,最终结果为:

css 复制代码
samples                  10000
loss                     0.411123
HBM Top-1                85.8900%
Quantized BC Top-1       85.8900%
差值                     0.0000pp
status                   PASS

HBM 与 Quantized BC 的精度完全一致,说明以下环节在 x86 和实板之间保持一致:

  • 图像预处理
  • INT8 输入量化
  • HBM 计算
  • INT32 输出反量化
  • Top-1 计算

13. 独立测试推理延时

完整精度测试包含数据读取、图像预处理和输入量化。延时测试使用一份提前完成预处理和量化的固定输入:

css 复制代码
预热帧数         100
正式测量帧数     1000
调用方式         单请求顺序调用
输入压缩         NONE
BPU              Core 0

测试时需要区分几种延时:

  • host_call:x86 Python 调用一次 session(...) 的总耗时。
  • sd2rv_duration:请求经 gRPC 发出到响应返回的耗时。
  • board_duration:板端预处理、推理和后处理总耗时。
  • infer_duration:BPU 纯推理耗时。

本次 BPU 纯推理平均耗时为 0.445 ms,板端总处理平均耗时为 0.633 ms,而 x86 侧看到的端到端平均耗时为 2.133 ms。

三个数字描述的范围不同。评估 BPU 性能时应关注 infer_duration;评估真实应用调用体验时,则需要关注端到端耗时。

相关推荐
IT星宿1 小时前
github-gravity:用 Rapier2D 把 GitHub 贡献图做成 SVG 物理动画
算法·github
♡すぎ♡2 小时前
SWE 的基本条件
算法·swe·浅水模拟
郝学胜-神的一滴3 小时前
AI 编程智能体 17:四步搞定专属AI智能体
大数据·人工智能·python·神经网络·算法·langchain·强化学习
微尘hjx3 小时前
【StrongSort 03】多轨迹算法StrongSort改进点
人工智能·算法·yolo·视觉检测·ultralytics·strongsort·轨迹算法
梦帮科技3 小时前
QML-04-QGraph-Diff:量子非欧图扩散生成模型与 SE(3) 等变分子晶格逆向合成完整技术指南(上部分)
人工智能·算法·自动化·量子计算·傅里叶分析
七牛开发者3 小时前
TokenGate:端到端可训练的 Token 剪枝机制,让大模型 Prefill 提速 11.3 倍
算法·机器学习·剪枝·长上下文·tokengate
牧瀬クリスだ5 小时前
Redis Hash实战:高效存储对象与购物车
java·数据库·redis·算法·哈希算法
木井巳5 小时前
【动态规划: 路径问题】下降路径最小和
java·算法·leetcode·动态规划·推荐算法