【X5】请问量化时如何避免_hz_cat_1_rescale_1这类操作的产生

开发与问题
,
11

抱歉，之前的代码示例有误！FakeQuantize 的构造函数不是直接传入 QConfig 对象。

错误原因

FakeQuantize 模块需要的是量化参数配置，而不是 QConfig 对象。QConfig 是用于 prepare_qat_fx 时配置每个 Module 的量化策略的。

正确方案：通过 prepare_qat_fx 的 qconfig 参数控制

要让 torch.cat 的多个输入共享量化参数，正确做法是在 prepare_qat_fx 时，为这些输入路径的子模块配置同一个 qconfig 实例，或者使用更高级的自定义 qconfig 映射。

方案1：为 cat 前的子模块设置共享 qconfig（推荐）

import torch
import torch.nn as nn
from horizon_plugin_pytorch.quantization import prepare_qat_fx, convert_fx
from horizon_plugin_pytorch.quantization.qconfig import default_qat_8bit_fake_quant_qconfig
import copy

class SharedCatModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(32, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
        self.conv3 = nn.Conv2d(32, 32, 3, padding=1)
        self.relu1 = nn.ReLU()
        self.relu2 = nn.ReLU()
        self.relu3 = nn.ReLU()
        
    def forward(self, x):
        x1 = self.relu1(self.conv1(x))
        x2 = self.relu2(self.conv2(x))
        x3 = self.relu3(self.conv3(x))
        return torch.cat([x1, x2, x3], dim=1)

model = SharedCatModule()

# 关键：创建一个共享的 qconfig 实例
shared_qconfig = default_qat_8bit_fake_quant_qconfig()

# 为 cat 的所有输入路径上的最后一个量化点设置相同的 qconfig
# 通常是 relu 或 conv 的输出
model.relu1.qconfig = shared_qconfig
model.relu2.qconfig = shared_qconfig
model.relu3.qconfig = shared_qconfig

# 准备量化
model = prepare_qat_fx(copy.deepcopy(model))

# calibration
calib_data = torch.rand(1, 32, 64, 64)
model.eval()
with torch.no_grad():
    model(calib_data)

# 转换
model = convert_fx(model)

方案2：使用 qconfig 字典映射（更灵活）

from horizon_plugin_pytorch.quantization import prepare_qat_fx
from horizon_plugin_pytorch.quantization.qconfig import default_qat_8bit_fake_quant_qconfig

# 创建共享的 qconfig 实例
shared_qconfig = default_qat_8bit_fake_quant_qconfig()

# 准备量化时，通过字典指定特定模块的 qconfig
model = prepare_qat_fx(
    copy.deepcopy(float_model),
    {
        "": default_qat_8bit_fake_quant_qconfig(),  # 默认配置
        "module_name": {
            # 为 cat 前的关键模块指定共享的 qconfig 实例
            "relu1": shared_qconfig,
            "relu2": shared_qconfig,
            "relu3": shared_qconfig,
        }
    }
)

方案3：使用 set_qconfig 方法（如果支持）

某些版本的 plugin 支持直接在模块上调用 set_qconfig：

# 在 prepare 之前
shared_qconfig = default_qat_8bit_fake_quant_qconfig()

for module in [model.relu1, model.relu2, model.relu3]:
    module.qconfig = shared_qconfig

model = prepare_qat_fx(copy.deepcopy(model))

验证共享是否生效

# calibration 后检查量化参数
def check_quant_params(model):
    print("=== 量化参数检查 ===")
    for name, module in model.named_modules():
        if hasattr(module, 'activation_post_process'):
            obs = module.activation_post_process
            if hasattr(obs, 'scale') and obs.scale is not None:
                print(f"{name}: scale={obs.scale.item():.6f}")

# 在 calibration 后调用
check_quant_params(model)

如果 relu1, relu2, relu3 的 scale 值相同，说明共享成功。

关键要点

要点 说明
:white_check_mark: QConfig 不是直接传给 FakeQuantize QConfig 是 prepare_qat_fx 的配置参数
:white_check_mark: 共享 qconfig 实例 多个模块的 .qconfig 属性指向同一个 QConfig 对象
:white_check_mark: 在 prepare_qat_fx 之前设置 prepare 时会读取模块的 qconfig 属性
:white_check_mark: 使用 copy.deepcopy 避免原模型被修改

如果仍有问题

请提供：

  1. pip show horizon-plugin-pytorch 的版本
  2. 你的模型中 cat 操作前后的具体结构
  3. 完整的 prepare_qat_fx 调用代码

建议参考官方文档：QAT 快速上手（fx mode）