地平线PrepareMethod：EAGER和Graph Model

地平线PrepareMethod 模式：

EAGER / SYMBOLIC

Eager Mode (急切模式)：
- 原理：像搭积木一样，你需要手动把每一个算子（Conv, ReLU, Add）都替换成支持量化的版本（如 QuantizedConv2d, FloatFunctional.add）。
- 痛点：对于包含循环、动态形状或复杂控制流（如 Deformable Attention）的模型，手动替换几乎是不可能的任务，且极易出错。
- 现状：官方已明确标记为"早期方案"，存在易用性问题，不建议使用。
SYMBOLIC (符号追踪)：
- 原理：早期的图追踪尝试，稳定性不如现在的 JIT/FX。
- 现状：同样被标记为"早期方案"，不建议使用。

Graph Mode (JIT / JIT_STRIP)

原理：
- 自动感知图结构 ：工具链通过 torch.jit.trace或类似机制，一次性跑通模型的前向传播，记录下完整的计算图。
- 自动替换与融合：基于记录下的图，工具链自动将浮点算子替换为量化算子，并自动执行算子融合（如 Conv+BN+ReLU）。
- 代码零侵入 ：你不需要修改模型内部的 forward逻辑，也不需要手动插入 QuantStub（除非用于界定边界）。

2. JIT vs JIT_STRIP：该如何选？

这两者底层技术一致，区别仅在于对"前后处理"的处理策略。

特性	`PrepareMethod.JIT`	`PrepareMethod.JIT_STRIP`
处理方式	对整个 traced graph进行量化处理。	识别 `QuantStub`和 `DeQuantStub`，跳过它们之外的部分。
适用场景	模型从头到尾都需要量化，或者没有明显的前后处理逻辑。	大多数实际场景。模型头部有预处理（如 Normalize），尾部有后处理（如 NMS, Softmax），这些部分通常不需要或不适合量化。
优势简单直接。	更精准。避免了对非核心逻辑（如数据预处理）的错误量化，减少了调试麻烦。
要求	无特殊要求。	必须在模型中正确插入 `QuantStub` (开头) 和 `DeQuantStub` (结尾)。

官方建议 ：**目前更推荐 PrepareMethod.JIT_STRIP。**因为它能自动剥离前后处理，适合作为大多数模型的起点。

3. 工程化落地建议

根据你的情况（使用官方 Plugin MultiScaleDeformableAttention），你应该采取以下策略：

步骤一：模型改造（仅添加 Stub）

在你的模型 forward函数中，确保输入经过 QuantStub，输出经过 DeQuantStub。

python 复制代码

import torch.nn as nn
from horizon_plugin_pytorch.quantization import QuantStub, DeQuantStub
from horizon_plugin_pytorch.nn import MultiScaleDeformableAttention

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = QuantStub() # 【新增】量化入口
        self.dequant = DeQuantStub() # 【新增】反量化出口
        
        self.deform_attn = MultiScaleDeformableAttention(...)
        # ...其他层...

    def forward(self, x):
        x = self.quant(x) # 【新增】标记量化开始
        
        # ...中间逻辑...
        x = self.deform_attn(x) 
        # ...中间逻辑...
        
        x = self.dequant(x) # 【新增】标记量化结束
        return x
    
    # 【删除】不要写 propagate_qconfig

步骤二：调用 Prepare（使用 JIT_STRIP）

python 复制代码

from horizon_plugin_pytorch.quantization import prepare, PrepareMethod, get_qconfig, QconfigSetter, ModuleNameTemplate, qint8, qint16

# 1.定义模板（外部控制精度）
templates = [
    ModuleNameTemplate({"": qint8}), #默认全INT8
    ModuleNameTemplate({"*.deform_attn": qint16}), #特殊层INT16
]

# 2.执行 Prepare
qat_model = prepare(
    model=float_model,
    example_inputs=example_input, # JIT模式必填，用于Trace图结构
    qconfig_setter=QconfigSetter(get_qconfig(), templates),
    method=PrepareMethod.JIT_STRIP, # 【关键】使用推荐的 JIT_STRIP
)

二、完整代码示例 (Eager Mode)

复制代码

class SplitedMSDA(BaseModule):
    def __init__(self, split_weight_mul=False, num_levels=4, init_cfg=None):
        super(SplitedMSDA, self).__init__(init_cfg=init_cfg)  # 此时 init_cfg 默认是 None
        self.split_weight_mul = split_weight_mul
        self.num_levels = num_levels
        self.sampling_mul3 = FloatFunctional()
        self.sampling_add2 = FloatFunctional()
        self.sampling_cat = FloatFunctional()
        self.attention_weight_sum = FloatFunctional()
        if self.split_weight_mul:
            self.attention_weight_mul = nn.ModuleList(
                (FloatFunctional() for _ in range(self.num_levels))
            )
            self.attention_weight_pre_sum = nn.ModuleList(
                (FloatFunctional() for _ in range(self.num_levels))
            )
        else:
            self.attention_weight_mul = FloatFunctional()

        print(" -> 【运行中】调用了：地平线拆分版 SplitedMultiScaleDeformableAttention")

    def _multi_scale_deformable_attn(
        self,
        value: Tensor,
        value_spatial_shapes: Tensor,
        sampling_locations: Tensor,
        attention_weights: Tensor,
    ) -> Tensor:
        """Fundamental implementation of multi-scale deformable attention.

        Args:
            value: The value has shape
                (bs, num_keys, num_heads, embed_dims//num_heads)
            value_spatial_shapes: Spatial shape of
                each feature map, has shape (num_levels, 2),
                last dimension 2 represent (h, w)
            sampling_locations: The location of sampling points,
                has shape
                (bs ,num_queries, num_heads, num_levels, num_points, 2),
                the last dimension 2 represent (x, y).
            attention_weights: The weight of sampling points
                used when calculate the attention, has shape
                (bs ,num_queries, num_heads, num_levels, num_points),

        Returns:
            Tensor: shape of (bs, num_queries, embed_dims)
        """
        # Convert Tensor to list to avoid multi cuda sync on enumerate.
        value_spatial_shapes = value_spatial_shapes.cpu().numpy().tolist()

        bs, _, num_heads, embed_dims = value.shape
        (
            _,
            num_queries,
            num_heads,
            num_levels,
            num_points,
            _,
        ) = sampling_locations.shape
        value_list = value.split(
            [int(H_ * W_) for H_, W_ in value_spatial_shapes],  # noqa: N806
            dim=1,
        )
        del value
        sampling_grids = self.sampling_mul3.mul_scalar(sampling_locations, 2)
        del sampling_locations
        sampling_grids = self.sampling_add2.add_scalar(sampling_grids, -1)
        sampling_value_list = []

        for level, (H_, W_) in enumerate(value_spatial_shapes):  # noqa: N806
            # bs, H_*W_, num_heads, embed_dims ->
            # bs, H_*W_, num_heads*embed_dims ->
            # bs, num_heads*embed_dims, H_*W_ ->
            # bs*num_heads, embed_dims, H_, W_
            value_l_ = (
                value_list[level]
                .reshape(bs, H_ * W_, -1)
                .transpose(1, 2)
                .reshape(bs * num_heads, embed_dims, H_, W_)
            )
            # bs, num_queries, num_heads, num_points, 2 ->
            # bs, num_heads, num_queries, num_points, 2 ->
            # bs*num_heads, num_queries, num_points, 2
            sampling_grid_l_ = (
                sampling_grids[:, :, :, level : (level + 1)]
                .transpose(1, 2)
                .reshape(bs * num_heads, num_queries, num_points, 2)
            )
            # bs*num_heads, embed_dims, num_queries, num_points
            sampling_value_l_ = F.grid_sample(
                value_l_,
                sampling_grid_l_,
                mode="bilinear",
                padding_mode="zeros",
                align_corners=False,
            )
            del value_l_
            del sampling_grid_l_

            if self.split_weight_mul:
                # (bs*num_heads, embed_dims, num_queries, num_points) *
                # (bs*num_heads, 1, num_queries, num_points) ->
                # (bs*num_heads, embed_dims, num_queries, num_points)
                sampling_value_l_ = self.attention_weight_mul[level].mul(
                    sampling_value_l_,
                    attention_weights[:, :, :, level, :]
                    .transpose(1, 2)
                    .reshape(bs * num_heads, 1, num_queries, num_points),
                )

                # (bs*num_heads, embed_dims, num_queries, num_points) ->
                # (bs*num_heads, embed_dims, num_queries, 1)
                sampling_value_l_ = self.attention_weight_pre_sum[level].sum(
                    sampling_value_l_, dim=-1, keepdim=True
                )

            sampling_value_list.append(sampling_value_l_)

        # if self.split_weight_mul:
        # (bs*num_heads, embed_dims, num_queries, 1) ->
        # (bs*num_heads, embed_dims, num_queries, num_levels)
        # else:
        # (bs*num_heads, embed_dims, num_queries, num_points) ->
        # (bs*num_heads, embed_dims, num_queries, num_levels*num_points)
        sampling_value_all = self.sampling_cat.cat(sampling_value_list, dim=-1)
        del sampling_value_list

        if not self.split_weight_mul:
            sampling_value_all = self.attention_weight_mul.mul(
                sampling_value_all,
                attention_weights.transpose(1, 2).reshape(
                    bs * num_heads, 1, num_queries, num_levels * num_points
                ),
            )

        # bs*num_heads, embed_dims, num_queries, 1
        output = self.attention_weight_sum.sum(
            sampling_value_all, dim=-1, keepdim=True
        ).view(bs, num_heads * embed_dims, num_queries)
        return output.transpose(1, 2)

    @typechecked
    def forward(
        self,
        value: Tensor,
        spatial_shapes: Tensor,
        sampling_locations: Tensor,
        attention_weights: Tensor,
    ) -> Tensor:
        output = self._multi_scale_deformable_attn(
                value, spatial_shapes, sampling_locations, attention_weights
            )
        return output

    def propagate_qconfig(self, qconfig):
        from horizon_plugin_pytorch.quantization.qconfig import (
            replace_int_activation_dtype,
        )

        replace_int_activation_dtype(self.sampling_mul3, qconfig, qint16)
        replace_int_activation_dtype(self.sampling_add2, qconfig, qint16)
        replace_int_activation_dtype(self.sampling_cat, qconfig, qint16)
        if self.split_weight_mul:
            for m in self.attention_weight_mul:
                replace_int_activation_dtype(m, qconfig, qint16)
            for m in self.attention_weight_pre_sum:
                replace_int_activation_dtype(m, qconfig, qint16)
        else:
            replace_int_activation_dtype(
                self.attention_weight_mul, qconfig, qint16
            )
        replace_int_activation_dtype(
            self.attention_weight_sum, qconfig, qint16
        )

        # 完整代码示例 (Eager Mode)
        from horizon_plugin_pytorch.quantization import prepare, PrepareMethod, get_qconfig
        from horizon_plugin_pytorch.quantization.qconfig_setter import QconfigSetter, default_calibration_qconfig_setter

        #1.定义全局默认 Setter (Eager模式下也需要 Setter来初始化根节点的 qconfig)  
        my_setter = QconfigSetter(    
            reference_qconfig=get_qconfig(),    
            templates=[default_calibration_qconfig_setter] #或其他模板  
        )  

        #2.执行 Prepare  
        float_model.eval()  
        example_input = (torch.randn(1, 64, H, W), torch.randn(1, N, 2), torch.randn(1, N))  

        qat_model = prepare(    
            float_model,    
            example_inputs=example_input,    
            qconfig_setter=my_setter,    
            method=PrepareMethod.EAGER, # <--- 必须指定    
        )  
        # 假设 deform_attn 是模型中的一个属性
        #print(qat_model.backbone.layer1.deform_attn.qconfig) 
        # 或者检查其内部算子的 dtype统计情况（在校准后）
        #3.校准  
        calibrate(qat_model, calib_data_loader)  

        #4.转换  
        quantized_model = convert(qat_model)

三、方案3

复制代码

        from horizon_plugin_pytorch.quantization import (
            prepare_qat_fx_based, 
            get_qconfig, 
            QconfigSetter, 
            ModuleNameTemplate, 
            qint8, 
            qint16
        )

        # 定义策略模板
        templates = [
            # 1. 【自动处理】全局默认 INT8
            # 这会覆盖模型中所有未被特殊指定的模块（包括 conv1等）
            ModuleNameTemplate({"": qint8}), 
            
            # 2. 【特殊处理】针对 Deformable Attention强制 INT16
            # 工具链会自动找到名为 deform_attn的模块并应用此配置
            # 其余部分依然保持 INT8，无需你手动遍历
            ModuleNameTemplate({
                "*.deform_attn": qint16 
            })
        ]

        # 执行 Prepare (FX模式)
        qat_model = prepare_qat_fx_based(
            float_model=my_net, 
            example_inputs=dummy_input, 
            qconfig_setter=QconfigSetter(get_qconfig(), templates)
        )