跟进建议
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 节点的方案,有结果随时同步 ![]()