PyTorch转ONNX：量化模型观察者名称与ONNX层名映射方法咨询

以下是几种可行的实现思路和代码示例：

1. 收集量化模型中观察者与PyTorch层的关联 #

量化后的PyTorch模型中，观察者（Observer）通常作为层的子模块存在（比如model.conv1.activation_post_process），或者本身就是独立模块。可以通过遍历模型模块，记录每个观察者对应的原始PyTorch层名：

import torch

def collect_observer_to_pytorch_layer(quantized_model):
    observer_layer_map = {}
    # 遍历模型所有模块
    for full_module_name, module in quantized_model.named_modules():
        # 处理带activation_post_process的层（静态量化常见情况）
        if hasattr(module, 'activation_post_process'):
            observer = module.activation_post_process
            # 生成唯一的观察者标识，比如类型+父层名
            observer_id = f"{observer.__class__.__name__}_{full_module_name.replace('.', '_')}"
            observer_layer_map[observer_id] = full_module_name
        # 处理本身就是观察者的模块（部分量化场景）
        elif isinstance(module, (torch.quantization.MinMaxObserver, 
                                torch.quantization.MovingAverageMinMaxObserver,
                                torch.quantization.PerChannelMinMaxObserver)):
            # 获取该观察者所属的父层名
            parent_layer_name = full_module_name.rsplit('.', 1)[0]
            observer_layer_map[full_module_name] = parent_layer_name
    return observer_layer_map

2. 导出ONNX并提取节点与PyTorch层的对应关系 #

PyTorch导出ONNX时，节点名称通常会保留原始PyTorch层的名称特征（可能会添加后缀如_0、_1）。可以通过加载ONNX模型遍历节点，或者导出时开启verbose=True查看节点信息：

import onnx

# 导出量化模型到ONNX（dummy_input为模型输入的样例张量）
torch.onnx.export(
    quantized_model,
    dummy_input,
    "quantized_model.onnx",
    opset_version=13,  # 建议使用适配量化的opset版本
    do_constant_folding=False,  # 避免常量折叠导致层名丢失
    verbose=True  # 开启后会打印所有ONNX节点信息，方便排查
)

# 加载导出的ONNX模型，提取节点名称与对应操作
def collect_onnx_node_mapping(onnx_model_path):
    onnx_model = onnx.load(onnx_model_path)
    onnx_node_map = {}
    for node in onnx_model.graph.node:
        # 记录ONNX节点名称和对应的操作类型
        onnx_node_map[node.name] = node.op_type
    return onnx_node_map

3. 建立观察者到ONNX层的最终映射 #

结合前两步的结果，通过PyTorch层名匹配ONNX节点名称（通常ONNX节点名会包含原PyTorch层名的关键词）：

# 获取两个基础映射
observer_to_pytorch = collect_observer_to_pytorch_layer(quantized_model)
onnx_node_map = collect_onnx_node_mapping("quantized_model.onnx")

# 构建最终的观察者-ONNX层映射
observer_to_onnx = {}
for observer_id, pytorch_layer_name in observer_to_pytorch.items():
    # 遍历ONNX节点，找到包含PyTorch层名的节点
    for onnx_node_name, op_type in onnx_node_map.items():
        if pytorch_layer_name in onnx_node_name:
            observer_to_onnx[observer_id] = onnx_node_name
            break

补充说明 #

• 如果是静态量化模型，导出ONNX时建议使用opset_version >= 13，更好支持量化节点的导出。
• 部分情况下，PyTorch层名中的特殊字符会被ONNX转义（比如.变成_），匹配时需要做对应处理。
• 若遇到名称匹配冲突，可以结合层的操作类型（比如Conv对应ONNX的Conv节点，Linear对应Gemm节点）进一步筛选。

内容的提问来源于stack exchange，提问作者חן רובין