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

跟进建议

Marcelo6151 提到的共享 QuantStub/FakeQuantize 节点是解决这个问题的关键。我补充一个可操作的代码示例:

核心方案:共享 FakeQuantize 节点

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
from horizon_plugin_pytorch.quantization import FakeQuantize

class CatModel(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.relu = nn.ReLU()
        
        # 关键:创建一个共享的 FakeQuantize 节点
        self.shared_quant = FakeQuantize(default_qat_8bit_fake_quant_qconfig())
        
    def forward(self, x):
        x1 = self.relu(self.conv1(x))
        x2 = self.relu(self.conv2(x))
        x3 = self.relu(self.conv3(x))
        
        # 所有输入通过同一个 FakeQuantize 节点,强制共享 scale/zero_point
        x1_q = self.shared_quant(x1)
        x2_q = self.shared_quant(x2)
        x3_q = self.shared_quant(x3)
        
        out = torch.cat([x1_q, x2_q, x3_q], dim=1)
        return out

model = CatModel()
model = prepare_qat_fx(model)

# calibration
calib_data = torch.rand(1, 32, 64, 64)
model(calib_data)

# 转换
model = convert_fx(model)

验证步骤

# 1. 检查量化参数是否真的共享了
def check_shared_quant(model):
    for name, module in model.named_modules():
        if isinstance(module, FakeQuantize):
            if hasattr(module, 'scale') and module.scale is not None:
                print(f"{name}: scale={module.scale}")

check_shared_quant(model)

# 2. 使用 perf_model 验证 rescale 节点是否减少
from horizon_plugin_pytorch import quantization as hq
hq.perf_model(
    model,
    [calib_data],
    opt=2,
    layer_details=True,
    out_dir="./perf_check"
)
# 查看输出日志中 hz_cat_1_rescale_1 的数量

如果问题仍未解决

请提供以下信息方便进一步排查:

信息项 说明
工具链版本 pip show horizon-plugin-pytorch
cat 操作位置 模型中 cat 前后的结构片段
perf_model 输出 hz_cat_1_rescale_1 的具体数量和位置
QAT 配置代码 完整的 prepare_qat_fx 调用代码

社区资源

先试试共享 FakeQuantize 节点的方案,有结果随时同步 :+1: