torchvision mobilenet_v2 导出 4w4f onnx 报错

[work-zhangzhe]

• 导出代码

import torchvision
import torch

from sparsebit.quantization import QuantModel, parse_qconfig


qconfig_path = "./qconfig_lsq.yaml"

# BACKEND: virtual
# W:
#   QSCHEME: per-channel-symmetric
#   QUANTIZER: 
#     TYPE: lsq
#     BIT: 4
# A:
#   QSCHEME: per-tensor-affine
#   QUANTIZER:
#     TYPE: lsq
#     BIT: 4
#   QADD:
#     ENABLE_QUANT: true

model = torchvision.models.mobilenet_v2(pretrained=True)
qconfig = parse_qconfig(qconfig_path)
qmodel = QuantModel(model, config=qconfig)
qmodel.eval()
inp = torch.randn(2, 3, 224, 224)
out = qmodel(inp)
print(out.shape)
with torch.no_grad():
    qmodel.export_onnx(
        inp, name="mobilenet_v2_4w4f.onnx", extra_info=True
    )

• 报错信息

Traceback (most recent call last):
  File "dump_onnx.py", line 17, in <module>
    qmodel.export_onnx(
  File "/data/Project/Sparsebit/sparsebit/quantization/quant_model.py", line 256, in export_onnx
    self.add_extra_info_to_onnx(name)
  File "/data/Project/Sparsebit/sparsebit/quantization/quant_model.py", line 294, in add_extra_info_to_onnx
    onnx_op = onnx_model.graph.node[op_pos]
IndexError: list index (914) out of range