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

开发与问题
,
19

错误原因分析

AssertionError: input scale must be the same as op's 这个错误的核心原因是：

你多次调用同一个 QuantStub 实例（如 self.quant_k），但每次输入的 tensor 数据分布不同，导致 calibration 后该 QuantStub 的 scale 无法同时满足所有输入路径的要求。

在 torch.jit.trace 时，工具链检测到 cat 的多个输入虽然来自同一个 QuantStub，但由于 forward 中实际传入的 tensor 不同（k_3x3, k_7x7, k_9x9），量化图认为这些输入应该有不同的 scale，但你又强制它们共享同一个 QuantStub，导致冲突。

正确解决方案

方案 1：为每个输入路径使用独立的 QuantStub，但强制共享 scale（推荐）

关键：在 prepare_qat_fx 后，手动绑定这些 QuantStub 的 observer，让它们共享同一组统计量。

import torch
import torch.nn as nn
import copy
from horizon_plugin_pytorch.quantization import (
    prepare_qat_fx,
    convert_fx,
    set_fake_quantize,
    FakeQuantState,
    QuantStub,
)
from horizon_plugin_pytorch.quantization.qconfig import default_qat_8bit_fake_quant_qconfig

class MultiScaleCatModule(nn.Module):
    def __init__(self, in_channels=128, window_sizes=[3, 7, 9]):
        super().__init__()
        self.window_sizes = window_sizes
        
        # 为每个输入路径创建独立的 QuantStub
        self.quant_k_3x3 = QuantStub()
        self.quant_k_7x7 = QuantStub()
        self.quant_k_9x9 = QuantStub()
        
        self.quant_v_3x3 = QuantStub()
        self.quant_v_7x7 = QuantStub()
        self.quant_v_9x9 = QuantStub()
        
        # 模拟不同感受野的特征提取（实际是你的卷积层）
        self.conv_3x3 = nn.Conv2d(in_channels, in_channels, 3, padding=1)
        self.conv_7x7 = nn.Conv2d(in_channels, in_channels, 7, padding=3)
        self.conv_9x9 = nn.Conv2d(in_channels, in_channels, 9, padding=4)
        
    def forward(self, x):
        # 提取多尺度特征
        k_3x3 = self.conv_3x3(x)
        k_7x7 = self.conv_7x7(x)
        k_9x9 = self.conv_9x9(x)
        
        v_3x3 = k_3x3  # 简化示例，实际可能不同
        v_7x7 = k_7x7
        v_9x9 = k_9x9
        
        # 量化后 cat
        k_features = torch.cat([
            self.quant_k_3x3(k_3x3),
            self.quant_k_7x7(k_7x7),
            self.quant_k_9x9(k_9x9),
        ], dim=1)
        
        v_features = torch.cat([
            self.quant_v_3x3(v_3x3),
            self.quant_v_7x7(v_7x7),
            self.quant_v_9x9(v_9x9),
        ], dim=1)
        
        return k_features, v_features

# ========== 关键步骤：共享 scale ==========

model = MultiScaleCatModule()

# 1. prepare
model = prepare_qat_fx(copy.deepcopy(model))

# 2. 【关键】手动绑定 K 路径的 observer，让它们共享统计量
def share_observers(model, stub_names):
    """强制多个 QuantStub 共享同一个 observer 实例"""
    if not stub_names:
        return
    
    # 获取第一个 stub 的 observer 作为主 observer
    first_stub = getattr(model, stub_names[0])
    if hasattr(first_stub, 'activation_post_process'):
        main_observer = first_stub.activation_post_process
        
        # 将其他 stub 的 observer 指向同一个实例
        for name in stub_names[1:]:
            stub = getattr(model, name)
            if hasattr(stub, 'activation_post_process'):
                stub.activation_post_process = main_observer

# 共享 K 路径的 scale
share_observers(model, ['quant_k_3x3', 'quant_k_7x7', 'quant_k_9x9'])

# 共享 V 路径的 scale
share_observers(model, ['quant_v_3x3', 'quant_v_7x7', 'quant_v_9x9'])

# 3. calibration
model.eval()
set_fake_quantize(model, FakeQuantState.CALIBRATION)

calib_data = torch.rand(1, 128, 64, 64)  # 示例数据
with torch.no_grad():
    # 多批数据校准
    for _ in range(10):
        model(calib_data)

# 4. 切换到 validation 状态
model.eval()
set_fake_quantize(model, FakeQuantState.VALIDATION)

# 5. 验证 scale 是否真的共享了
def check_scales(model, stub_names):
    scales = []
    for name in stub_names:
        stub = getattr(model, name)
        if hasattr(stub, 'activation_post_process'):
            obs = stub.activation_post_process
            if hasattr(obs, 'scale') and obs.scale is not None:
                scales.append(obs.scale.item())
    print(f"{stub_names} scales: {scales}")
    if len(set(scales)) == 1:
        print("✅ Scale 共享成功！")
    else:
        print("❌ Scale 未共享，需要检查绑定逻辑")

check_scales(model, ['quant_k_3x3', 'quant_k_7x7', 'quant_k_9x9'])
check_scales(model, ['quant_v_3x3', 'quant_v_7x7', 'quant_v_9x9'])

# 6. 转换
model = convert_fx(model)

# 7. trace 并编译
example_input = torch.rand(1, 128, 64, 64)
traced_model = torch.jit.trace(model.cpu(), example_input)

# 8. 性能验证
from horizon_plugin_pytorch import quantization as hq
hq.perf_model(
    traced_model,
    [example_input],
    opt=2,
    layer_details=True,
    out_dir="./perf_check"
)

方案 2：使用 torch.ao.quantization 的 share_qconfig（如果版本支持）

某些版本的 plugin 支持在 qconfig_dict 中使用特殊语法，但需要确认你的版本是否支持。

方案 3：如果上述方案因版本问题不可行，改用 PTQ

PTQ 的 YAML 配置可能更灵活，支持 share_observer 参数：

# 量化配置 YAML
compiler_parameters:
  share_observer: true
  # 或针对特定 layer
  layer_parameters:
    - layer_name: "quant_k_3x3"
      share_observer_with: "quant_k_7x7"
    - layer_name: "quant_k_7x7"
      share_observer_with: "quant_k_9x9"

关键要点

要点 说明
:cross_mark: 不要在 forward 中多次调用同一个 QuantStub 会导致 scale 冲突
:white_check_mark: 为每个输入路径创建独立的 QuantStub 避免 trace 时的 scale 检查失败
:white_check_mark: 在 prepare 后手动绑定 observer 强制共享统计量
:white_check_mark: calibration 需要多批数据 确保统计量准确

如果仍有问题

请提供：

  1. pip show horizon-plugin-pytorch 的版本
  2. 完整的模型代码（包括 k_3x3, k_7x7, k_9x9 的来源）
  3. perf_model 输出的详细日志

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

如果问题复杂，建议联系 FAE 支持（如 Marcelo6151 所说）:+1: