Prepare说明

在 PyTorch 模型量化基本流程说明、构建浮点模型和构建 QConfig 章节中，我们已经介绍了 Prepare 在整条链路中的位置。这一阶段的目标，是把浮点模型转换为可用于 Calibration 或 QAT 的伪量化模型，并生成一组检查产物，帮助您确认图结构、量化范围和 QConfig 是否基本符合预期。

Prepare 阶段最常见的输入和输出如下：

项目	说明
输入	`float model`、`example_inputs`（或 `example_kw_inputs`）、qconfig 配置
主要输出	`calib model` 或 `qat model`
关键检查产物	`model_check_result.txt`：模型结构检查结果 `fx_graph.txt`：FX 图的文本描述（tabular 格式） `fx_graph.onnx`：FX 图的 ONNX 可视化文件 `jit_graph.pt`：JIT 序列化图，可作为后续 `reference_graph` 的基准 `graph_compare.html`：图对比报告（设置 `reference_graph` 时生成） `qconfig_dtypes.pt`、`qconfig_dtypes.pt.py`：量化配置 dtype 信息 `qconfig_changelogs.txt`：量化配置变更记录
下一步	如果检查结果正常，通常继续进入 Calibration

Prepare 做了什么

Prepare 一般会完成下面几件事：

计算图捕获。执行模型完整的 forward 逻辑，并在此过程中捕获对应的计算图。后续会根据 QuantStub 和 DeQuantStub 的位置裁剪图结构，只保留真正需要量化的部分。
function 类算子替换。部分 torch function 类型算子在量化时需要插入伪量化节点，因此会按图结构替换成对应的 Module 类型实现。不在图中的操作不会被替换。
算子融合。对满足条件的 pattern 执行融合，避免中间结果被额外量化。相关原理可参考算子融合。
算子转换。将浮点算子替换为 qat 算子，并按照 qconfig 在输入、输出或权重位置插入伪量化或伪转换节点。
模型结构检查。生成 model_check_result.txt、fx_graph.txt 等检查产物，供后续确认模型结构、图捕获结果和 dtype 配置是否正常。

注意

请确保 Prepare 之后不会再修改模型。否则已经被替换的 qat 算子可能产生不符合预期的行为，qconfig 应用结果、图结构以及后续导出行为也都可能与预期不一致。

例如，Prepare 之后再把未融合的 BN 改成 SyncBN，可能导致 qat BN 被再次修改为 SyncBN。类似的结构性改动，应当放在 Prepare 之前完成。

注解

Prepare 后的模型在数值上和浮点模型并不完全相同。主要原因有两类：

模型中已经插入了伪量化节点。
极少数输出存在极大值的算子，例如 reciprocal，为了适配量化，输出会被默认裁剪到更合理的范围。

因此，Prepare 后看到少量数值差异是正常现象，更重要的是先确认图结构和量化配置是否正确。

基本用法

如果您使用的是 QconfigSetter + templates 这条主线，通常推荐结合图模式使用 Prepare。

prepare 接口的完整签名如下：

def prepare(
    model: torch.nn.Module,
    example_inputs: Any = None,
    qconfig_setter: Optional[ModernQconfigSetterBase] = None,
    method: PrepareMethod = PrepareMethod.JIT_STRIP,
    example_kw_inputs: Any = None,
    check_result_dir: Optional[str] = None,
    inplace: bool = False,
    reference_graph: Optional[str] = None,
) -> torch.nn.Module:

各参数说明如下：

参数	类型	说明
`model`	`torch.nn.Module`	需要 Prepare 的浮点模型
`example_inputs`	`Any`	模型输入，用于 trace 和检查模型。当 `method` 为 `JIT_STRIP` 或 `JIT` 时必填
`qconfig_setter`	`Optional[ModernQconfigSetterBase]`	qconfig 配置，支持 `QconfigSetter` 模版
`method`	`PrepareMethod`	Prepare 模式，默认为 `PrepareMethod.JIT_STRIP`
`example_kw_inputs`	`Any`	模型 kwargs 输入，用于 trace 和检查模型，必须为 dict 类型
`check_result_dir`	`Optional[str]`	保存模型检查结果（qat check result）的目录路径
`inplace`	`bool`	是否原地修改模型，默认为 `False`（复制一份后再 Prepare）
`reference_graph`	`Optional[str]`	用于与本次 Prepare 结果进行图对比的参考图。支持本地 `.pt` 文件路径（推荐使用上次 Prepare 在 `check_result_dir` 下自动保存的 `jit_graph.pt`）、以 `http://` 或 `https://` 开头的集群 `.pt` 文件链接。仅在 `method` 为 `JIT_STRIP` 或 `JIT` 时支持。不为 `None` 时，会在 `check_result_dir`（或 `./.model_check_results`）下生成可视化对比报告 `graph_compare.html`，详见图对比功能

基本用法示例：

from horizon_plugin_pytorch.quantization import get_qconfig, prepare, PrepareMethod, qint8
from horizon_plugin_pytorch.quantization.qconfig_setter import (
    ConvDtypeTemplate,
    ModuleNameTemplate,
    QconfigSetter,
    SensitivityTemplate,
)

qat_model = prepare(
    float_model,
    example_inputs=example_inputs,
    qconfig_setter=QconfigSetter(
        reference_qconfig=get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
            SensitivityTemplate(table, ratio=0.2),
        ],
    ),
    method=PrepareMethod.JIT_STRIP,
)

这里建议您重点注意三点：

example_inputs 需要能正常跑通 forward，它决定了 Prepare 实际捕获到的图结构。同时 method 为 JIT_STRIP 时也必须提供 example_inputs（或 example_kw_inputs）。
qconfig_setter 会把模板配置应用到图中的算子上，因此最终结果应以 Prepare 产物中的实际 dtype 为准，而不是只看配置代码本身。
Prepare 完成后，建议立刻检查 model_check_result.txt 和 fx_graph.txt，不要等到 Calibration 或 QAT 结果异常后再回头补看。

PrepareMethod

prepare method 包括 JIT_STRIP 和 EAGER。前者属于 Graph Mode，后者属于 Eager Mode。它们的差异如下：

method	原理	优点	风险或代价
`PrepareMethod.JIT_STRIP`	使用 hook 和 subclass tensor 的方式感知图结构，在原有 forward 上做算子替换、算子融合等操作	自动化程度高，代码改动少，便于先快速跑通主链路	动态代码块、分支路径和循环次数变化需要重点检查
`PrepareMethod.EAGER`	不感知图结构，算子替换和算子融合主要依赖手工处理	灵活，可控，适合特殊需求	手动操作较多，接入和维护成本更高

目前更推荐 PrepareMethod.JIT_STRIP。它会根据模型中 QuantStub 和 DeQuantStub 的位置识别并跳过前后处理，因此更适合作为大多数模型的起点。

图模式使用示例

下面的简化示例主要帮助您理解 JIT_STRIP 下的几个典型行为：

import copy
import numpy as np
import torch
from torch.nn import functional as F
from torch.quantization import DeQuantStub, QuantStub

from horizon_plugin_pytorch.fx.jit_scheme import Tracer
from horizon_plugin_pytorch.quantization import PrepareMethod, prepare

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.quant0 = QuantStub()
        self.quant1 = QuantStub()
        self.dequant = DeQuantStub()
        ...

    def forward(self, input, other, target=None):
        # 不需要量化的前处理，使用 JIT_STRIP 时会从计算图中剔除
        input = (input - 128) / 128.0

        x = self.quant0(input)
        y = self.quant1(other)

        n = np.random.randint(1, 5)
        m = np.random.randint(1, 5)
        for _ in range(n):
            for _ in range(m):
                with Tracer.dynamic_block(self, "ConvBnAdd"):
                    x = self.conv(x)
                    x = self.bn(x)
                    x = x + y

        x = self.classifier(x).squeeze()
        if self.training:
            x = self.dequant(x)
            return F.cross_entropy(torch.softmax(x, dim=1), target)
        return torch.argmax(x, dim=1)

model.eval()
qat_model = prepare(
    model,
    example_inputs=copy.deepcopy(eval_example_input),
    method=PrepareMethod.JIT_STRIP,
)

qat_model.graph.print_tabular()

在这个例子里，input = (input - 128) / 128.0 这类不需要部署的前处理会被剥离出图。动态循环本身不会被改写，真正需要标注的是循环内部涉及算子替换或算子融合的代码块。如果您打印 qat_model.graph，其中的 scope_end 是 trace 过程中自动插入的边界标记，用于标记子 module 或动态代码块的边界，不对应实际计算。

注意

动态代码块涉及到算子替换或算子融合时，必须使用 Tracer.dynamic_block 进行标注，否则容易导致量化信息错乱或 forward 报错。
模型中调用次数变化的部分，例如某个子 module 或某个 dynamic_block，如果在 trace 时仅执行了一次，则有可能和非动态部分产生算子融合，导致 forward 报错。

开始前先确定平台起点

Prepare 本身负责把模型转换成伪量化模型，但在真正执行之前，通常需要先把当前平台更适合的 qconfig 起点确定下来。这样做的目的，不是一次就把最终精度配置定死，而是先用更符合平台能力的基础配置完成第一次 Prepare 和后续 Calibration，再根据结果继续做局部调整。

如果您还没有决定从哪类精度配置开始，可以先按平台区别做下面的选择。

J6E/M/B 平台常见起点

J6E/M/B 平台以定点算力为主，更常见的起点通常是全局 int8，再将部分量化敏感的算子提升到 int16。

使用 SensitivityTemplate 自动将 topk 或一定比例的敏感算子提升为 int16，示例如下：

table1 = torch.load("output1_ATOL_sensitive_ops.pt")
table2 = torch.load("output2_ATOL_sensitive_ops.pt")

my_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(observer=HistogramObserver),
    templates=[
        ModuleNameTemplate({"": qint8}),
        ConvDtypeTemplate(input_dtype=qint8, weight_dtype=qint8),
        MatmulDtypeTemplate(input_dtypes=qint8),
        SensitivityTemplate(
            sensitive_table=table1,
            topk_or_ratio=10,
        ),
        SensitivityTemplate(
            sensitive_table=table2,
            topk_or_ratio=0.1,
        ),
    ],
)

topk_or_ratio 的选择需要结合量化精度和部署性能一起权衡。一般来说，高精度算子越多，量化精度越好，部署性能影响也会越大。

J6P/H 平台常见起点

J6P/H 平台具备浮点算力，更常见的基础配置通常是全局 fp16，同时将 Conv / Matmul 配置为 int8，然后再将量化敏感的算子提升到 int16。

例如：

my_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(observer=HistogramObserver),
    templates=[
        ModuleNameTemplate({"": torch.float16}),
        ConvDtypeTemplate(input_dtype=qint8, weight_dtype=qint8),
        MatmulDtypeTemplate(input_dtypes=[qint8, qint8]),
        ConvDtypeTemplate(
            input_dtype=qint8,
            weight_dtype=qint16,
            prefix=["backbone.conv1.conv1_3.conv"],
        ),
        MatmulDtypeTemplate(
            input_dtypes=[qint16, qint8],
            prefix=["encoder.encoder.0.layers.0.self_attn.matmul"],
        ),
    ],
)

如果您已经有敏感度分析结果，也可以继续使用 SensitivityTemplate 做批量提升，例如：

table1 = torch.load("output1_ATOL_sensitive_ops.pt")
table2 = torch.load("output2_ATOL_sensitive_ops.pt")

my_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(observer=HistogramObserver),
    templates=[
        ModuleNameTemplate({"": torch.float16}),
        ConvDtypeTemplate(input_dtype=qint8, weight_dtype=qint8),
        MatmulDtypeTemplate(input_dtypes=[qint8, qint8]),
        SensitivityTemplate(
            sensitive_table=table1,
            topk_or_ratio=10,
        ),
        SensitivityTemplate(
            sensitive_table=table2,
            topk_or_ratio=0.1,
        ),
    ],
)

这些起点更适合帮助您先完成第一次 Prepare 和 Calibration。如果后续还需要继续调整模板细节、理解不同模板的覆盖关系或处理 fix scale，可再回到构建 QConfig 章节继续展开。

Prepare 成功后建议检查什么

Prepare 完成后，建议至少先确认下面几件事：

模型能否正常完成一次 forward。
model_check_result.txt 中是否存在明显异常提示。
fx_graph.txt 中的图结构是否符合预期。
qconfig_dtypes.pt.py 中的关键模块 dtype 是否基本符合您的模板配置。
qconfig_changelogs.txt 中是否出现了您未预期的覆盖或回退。

先看 `model_check_result.txt`

这个文件适合在 Prepare 刚完成后做第一次总检查。建议重点关注下面几类信息：

在提供 example_inputs 的情况下，Prepare 默认会对模型结构做一次检查，并生成 model_check_result.txt。如果检查失败，通常需要先根据 warning 修改模型写法或量化配置，必要时再单独调用 check_qat_model 继续检查。

是否仍存在未融合的 pattern。如果这里还能看到本应融合的 conv / bn / add / relu 等 pattern，通常说明当前模型结构或写法没有满足预期融合条件。
是否存在共享模块。若某个模块被多次调用，需要进一步判断这些调用的输出分布是否接近。若差异较大，可能继续影响 scale 对齐、训练行为或后续导出一致性。
是否出现异常 dtype。重点确认各算子的 input / weight / output 配置是否与预期一致，是否存在工具自动回退，或是否存在工具为提升量化精度而自动保留的高精度计算。
是否出现异常 qconfig 提示。例如 fixed scale、模板覆盖结果或某些层的特殊配置是否真的符合预期，都应在这里先确认一遍。

如果您现在还看不懂这个文件里每一类结果分别代表什么，建议继续阅读调试产物解读。

再看 `fx_graph.txt`

这个文件更适合确认工具最终捕获到的执行路径和图结构，而不是直接判断“哪里错了”。建议重点看下面几件事：

实际被捕获的执行路径是否就是部署时真正使用的路径。
QuantStub、DeQuantStub、算子替换和量化节点插入的位置是否符合预期。
动态代码块边界是否正确，某些应当进入图的逻辑是否被遗漏，某些不该进入量化范围的逻辑是否被错误带入。
在 init 中定义的模块，是否真的在本次 Prepare 的 forward 路径中被调用。

如果 fx_graph.txt 本身就已经和预期结构明显不一致，通常不建议继续做后面的 Calibration 或精度调优，而应先回到 Prepare 阶段调整模型写法、example_inputs 或分支边界。

图对比功能

Prepare 提供图对比功能，可通过 reference_graph 参数将本次 Prepare 生成的图与一份历史参考图进行结构化对比，并自动生成可读 HTML 报告，帮助快速识别图结构差异。

适用场景：不同阶段（train/eval/deploy）、代码重构、或模型结构调整后，确认两次 Prepare 的图结构是否一致。

支持的 method：仅在 method 为 PrepareMethod.JIT_STRIP 时支持。Eager 模式无 graph 生成，不支持此功能。

reference_graph 参数格式：

本地 .pt 文件路径（推荐使用上次 Prepare 在 check_result_dir 下自动保存的 jit_graph.pt）。
以 http:// 或 https:// 开头的集群链接。

使用流程：

第一次 prepare，按 deploy 逻辑进行 prepare，reference_graph 参数保持默认 None，得到 deploy graph jit_graph.pt，将此作为基准。建议您将 jit_graph.pt 重命名，如 base_graph.pt，防止后续使用相同路径时被覆盖。
第二次 prepare，在 calib/QAT/train/eval 状态下，将 deploy graph 作为基准进行对比。reference_graph 指向上次的 base_graph.pt（或集群链接）。prepare 会生成 graph 对比报告 graph_compare.html，您可以使用浏览器打开查看。

示例：

import torch
from horizon_plugin_pytorch import March, set_march
from horizon_plugin_pytorch.quantization import (
    prepare,
    get_qconfig,
)
from horizon_plugin_pytorch.quantization.qconfig_setter import *

set_march(March.NASH_E)
example = torch.randn(1, 3, 224, 224)

# 步骤 1：第一次 prepare，保存基准图，以 deploy 流程为准
out_dir = "./qat_prepare_runs/run_v1"
float_model_v1.deploy = True
qat_v1 = prepare(
    float_model_v1,
    example_inputs=example,
    qconfig_setter=QconfigSetter(
        reference_qconfig=get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
            MatmulDtypeTemplate(),
        ],
    ),
    check_result_dir=out_dir,
)

# 步骤 2：第二次 prepare，将 deploy 时的 graph 作为基准进行对比
reference_path = f"{out_dir}/jit_graph.pt"  # 本地路径
# 也可配置集群路径
# reference_path = "https://your-cluster.example.com/path/to/run_v1/jit_graph.pt"

float_model_v1.deploy = False
out_dirv2 = "./qat_prepare_runs/run_v2"
qat_v2 = prepare(
    float_model_v1,
    example_inputs=example,
    qconfig_setter=QconfigSetter(
        reference_qconfig=get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
            MatmulDtypeTemplate(),
        ],
    ),
    check_result_dir=out_dirv2,
    reference_graph=reference_path,
)

在完成上述两个步骤后，您就可以在浏览器内打开对比报告 graph_compare.html 进行查看。

对比报告说明：

当 reference_graph 不为 None 时，Prepare 会在 check_result_dir 目录（若未设置则为 ./.model_check_results）下生成 graph_compare.html 文件。该报告为左右并排表格，如下所示：

左侧：参考图（model 1 / old）。
右侧：本次 Prepare 图（model 2 / new）。
- 红底：仅在参考图中存在的节点。
+ 绿底：仅在本次图中存在的节点。
空格 / 白底：两侧节点对齐且字段一致。
连续多行完全相同：默认折叠，点击蓝色摘要栏可展开。

运行日志中会输出报告路径，例如：JIT graph compare HTML written to .../graph_compare.html。

提示

如果模型结构完全相同，对比报告会显示 identical: True；若模型不同，则显示 identical: False 并标出差异的详细位置。

需要重点关注的风险

动态代码块、deploy 分支和循环次数变化

下面几类写法，最容易让图捕获结果和预期不一致：

训练和部署阶段走不同分支，例如 if self.training 或 if deploy。
forward 中存在依赖条件判断的动态路径。
循环次数在不同场景下不一致。
某些模块在 init 中定义了，但当前 Prepare 的 forward 路径中并未实际调用。

出现这些情况时，最重要的不是先猜是哪一步错了，而是先结合 fx_graph.txt 确认工具到底按什么图在处理模型。

如果动态代码块中涉及到算子替换或算子融合，必须使用 Tracer.dynamic_block 之类的方式显式标注。否则很容易出现量化信息错乱、融合结果异常，甚至 forward 报错。

一个典型写法如下：

from horizon_plugin_pytorch.fx.jit_scheme import Tracer

for _ in range(n):
    for _ in range(m):
        with Tracer.dynamic_block(self, "ConvBnAdd"):
            x = self.conv(x)
            x = self.bn(x)
            x = x + y

这里真正需要标注的，是动态循环中涉及算子替换或算子融合的代码块，而不是 for 循环本身。

function 算子替换与 Scope

function 算子替换依赖计算图。图中同一行代码的多次调用会展开成多个节点，因此也可能被替换成多个 Module。如果模型中存在训练和推理阶段调用次数不一致的代码块，就可能产生 scale 错位并进一步影响模型精度。

为了解决这个问题，原始实现里引入了 Scope 的概念。同一个 Scope 中，同一次 func 代码调用会替换为同一个共享 Module。常见的 Scope 定义方式有两类：

一个 Module 类型的 forward 方法。
使用 dynamic_block 标记的代码块。

除此以外，Prepare 还会通过 hook 完成部分 Wrapper 的封装和解封。因此在 Prepare 完成后，请不要修改模型中的任何 hook，以免替换失效，造成报错或精度问题。

尽量把部署逻辑单独剥离出来

使用基于图的 Prepare 方法时，最好保证进入 Prepare 的 forward 只包含部署逻辑。如果训练分支、loss 计算或其他非部署逻辑直接混在一起，就更容易让图捕获结果和最终部署逻辑不一致。

下面是一个更稳妥的写法示例：

# 只对部署逻辑做 prepare
def forward_infer(self, input, gt):
    conv_out = self.conv(input)
    return self.sigmoid(conv_out), conv_out

# 非部署逻辑留在外层
def forward(self, input, gt):
    sig_out, conv_out = self.forward_infer(input, gt)
    if self.training:
        return self.loss(conv_out, gt)
    return sig_out

如果代码因为可读性或维护性原因，暂时无法剥离出干净的部署逻辑，建议至少额外做两项检查：

检查加载 ckpt 时是否存在 missing key 和 unexpected key。若量化参数缺失或多出，往往意味着本次 Prepare 的 forward 路径和保存 ckpt 时的路径没有对齐。
比较多次 Prepare 生成的 fx_graph.txt 是否一致。若多次生成的图结构本身就不一致，通常说明 forward 路径存在不稳定因素，需要先判断这是否符合预期。

异常处理

如果您在 Prepare 阶段已经发现异常，通常可以按下面的方向继续：

如果暂时看不懂 model_check_result.txt、fx_graph.txt 或 dtype 检查结果，建议先看调试产物解读。
如果怀疑 Prepare 后图结构变化、导出图不一致，或训练和部署路径本身就不一致，建议继续看部署一致性分析说明。
如果问题更像是模板配置、dtype 配置或高精度输出策略本身不符合预期，建议回到构建 QConfig 继续调整。
如果 Prepare 阶段的检查结果基本正常，下一步通常继续进入 Calibration。

#Prepare说明

#Prepare 做了什么

#基本用法

#PrepareMethod

#图模式使用示例

#开始前先确定平台起点

#J6E/M/B 平台常见起点

#J6P/H 平台常见起点

#Prepare 成功后建议检查什么

#先看 model_check_result.txt

#再看 fx_graph.txt

#图对比功能

#需要重点关注的风险

#动态代码块、deploy 分支和循环次数变化

#function 算子替换与 Scope

#尽量把部署逻辑单独剥离出来

#异常处理