j6-plugin-export

为 Horizon QAT 模型创建独立 HBIR 导出脚本（hbdk4.export）

目标

在使用 horizon_plugin_pytorch 的 QAT 工具链完成校准/训练后，创建一个独立的导出脚本，将 QAT 模型导出为 HBIR IR 模块，用于后续编译部署。

本 Skill 只做一件事：创建独立导出脚本，脚本中包含

from horizon_plugin_pytorch.quantization import hbdk4 as hb4
from hbdk4.compiler import save

model.eval()
set_fake_quantize(model, FakeQuantState.VALIDATION)
hbir_model = hb4.export(model, example_inputs, ...)
save(hbir_model, "output_path")

不引入任何其他改动（不改模型结构、不改 qconfig、不改 prepare/convert、不改训练/数据逻辑、不在训练或评测脚本中添加导出逻辑）。

强约束（必须遵守）

1) 必须以独立导出脚本的形式执行

导出逻辑必须放在独立脚本中，禁止在训练脚本、评测脚本或其他已有脚本中添加导出逻辑。

原因：

• 导出是一次性动作，不应与训练/评测流程耦合
• 独立脚本便于单独调试、复用和版本管理
• 避免训练/评测脚本因导出逻辑变得臃肿

2) 独立导出脚本的标准结构

独立导出脚本必须包含以下步骤，按固定顺序执行：

1. set_march — 设置目标平台
2. 构建/加载 QAT 模型 — 从 checkpoint 恢复
3. model.eval() — 切换到推理模式
4. set_fake_quantize(model, FakeQuantState.VALIDATION) — 切换到验证状态
5. QAT 模型验证 — 检查模型包含 fake-quant 模块，确保是 QAT 模型
6. 构造 example_inputs
7. hb4.export(model, example_inputs, ...) — 执行导出
8. save(hbir_model, output_path) — 保存 HBIR 模型到文件

3) 不在训练/评测脚本中添加导出逻辑

即使训练/评测脚本末尾是"导出的自然位置"，也不应在其中插入导出代码。应创建独立脚本。

4) 导出前必须验证模型是 QAT 模型

hb4.export 不区分 QAT 模型和 float 模型，如果误传 float 模型，导出会静默成功但结果无意义。因此：

Agent 侧约束：创建导出脚本前，必须确认用户提供的 checkpoint 是 QAT checkpoint（已完成 prepare + 校准/训练），而非 float checkpoint。如果用户未明确说明，应主动询问确认。

运行时检查：导出脚本中必须在 export 之前加入 QAT 模型验证，检查模型是否包含 fake-quant 模块：

# 验证模型是 QAT 模型（包含 fake-quant 模块）
_has_fq = any(
    "FakeQuantize" in type(m).__name__
    for m in model.modules()
)
assert _has_fq, (
    "当前模型不包含 FakeQuantize 模块，不是 QAT 模型。"
    "请确认使用的是 QAT checkpoint 而非 float checkpoint。"
)

标准改法（通用模板）

1) 增加 import

import torch
import horizon_plugin_pytorch as horizon
from horizon_plugin_pytorch.quantization import hbdk4 as hb4
from horizon_plugin_pytorch.quantization import FakeQuantState, set_fake_quantize
from hbdk4.compiler import save

2) 独立导出脚本完整模板

import torch
import horizon_plugin_pytorch as horizon
from horizon_plugin_pytorch.quantization import hbdk4 as hb4
from horizon_plugin_pytorch.quantization import FakeQuantState, set_fake_quantize
from hbdk4.compiler import save

# 1. 设置 march
horizon.march.set_march(horizon.march.March.NASH_E)

# 2. 构建/加载 QAT 模型
model = build_model()
model.load_state_dict(torch.load("qat_checkpoint.pth"))

# 3. eval
model.eval()

# 4. 切换到 VALIDATION
set_fake_quantize(model, FakeQuantState.VALIDATION)

# 5. 验证是 QAT 模型
_has_fq = any("FakeQuantize" in type(m).__name__ for m in model.modules())
assert _has_fq, (
    "当前模型不包含 FakeQuantize 模块，不是 QAT 模型。"
    "请确认使用的是 QAT checkpoint 而非 float checkpoint。"
)

# 6. 构造 example_inputs
example_input = torch.randn(1, 3, 224, 224)

# 7. 导出
hbir_model = hb4.export(model, (example_input,))

# 8. 保存
save(hbir_model, "output.bc")