Errors when convert quantized model to onnx

[qiuzhewei]

Hello there,
I am trying QAT my model and convert the quantized model to onnx referring to https://github.com/pytorch/ao/blob/main/torchao/quantization/qat/README.md. During the onnx conversion, an error occurs, blow is the detailed information:

/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2817:0
  %22461 : Float(1, 960, 7, 7, strides=[47040, 49, 7, 1], requires_grad=1, device=cpu) = aten::relu(%input.297), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.container.Sequential::blocks/torch.nn.modules.container.Sequential::blocks.4/timm.models._efficientnet_blocks.ConvBnAct::blocks.4.0/timm.layers.norm_act.BatchNormAct2d::bn1/torch.nn.modules.activation.ReLU::act # /usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:1699:0
  %input.301 : Float(1, 960, 1, 1, strides=[960, 1, 1, 1], requires_grad=1, device=cpu) = aten::adaptive_avg_pool2d(%22461, %21989), scope: timm.models.mobilenetv3.MobileNetV3::/timm.layers.adaptive_avgmax_pool.SelectAdaptivePool2d::global_pool/torch.nn.modules.pooling.AdaptiveAvgPool2d::pool # /usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:1379:0
  %x.203 : Float(1, 1280, 1, 1, strides=[1280, 1, 1, 1], requires_grad=0, device=cpu) = aten::_convolution(%input.301, %conv_head.weight, %18935, %21989, %21991, %21989, %22463, %21991, %22464, %22465, %22463, %22465, %22465), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.conv.Conv2d::conv_head # /usr/local/lib/python3.12/dist-packages/torch/nn/modules/conv.py:543:0
  %input : Float(1, 1280, 1, 1, strides=[1280, 1, 1, 1], requires_grad=1, device=cpu) = aten::batch_norm(%x.203, %norm_head.weight, %norm_head.bias, %norm_head.running_mean, %norm_head.running_var, %22463, %22466, %22467, %22465), scope: timm.models.mobilenetv3.MobileNetV3::/timm.layers.norm_act.BatchNormAct2d::norm_head # /usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:2817:0
  %22462 : Float(1, 1280, 1, 1, strides=[1280, 1, 1, 1], requires_grad=1, device=cpu) = aten::relu(%input), scope: timm.models.mobilenetv3.MobileNetV3::/timm.layers.norm_act.BatchNormAct2d::norm_head/torch.nn.modules.activation.ReLU::act # /usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:1699:0
  %22475 : Long(device=cpu) = prim::Constant[value={-1}](), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.flatten.Flatten::flatten
  %x : Float(1, 1280, strides=[1280, 1], requires_grad=1, device=cpu) = aten::flatten(%22462, %22464, %22475), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.flatten.Flatten::flatten # /usr/local/lib/python3.12/dist-packages/torch/nn/modules/flatten.py:53:0
  %22476 : Double(device=cpu) = prim::Constant[value={0.2}](), scope: timm.models.mobilenetv3.MobileNetV3::
  %input_tensor : Float(1, 1280, strides=[1280, 1], requires_grad=1, device=cpu) = aten::dropout(%x, %22476, %22463), scope: timm.models.mobilenetv3.MobileNetV3:: # /usr/local/lib/python3.12/dist-packages/torch/nn/functional.py:1422:0
  %21943 : str = prim::Constant[value="ASYMMETRIC"](), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.linear.Linear::classifier # /usr/local/lib/python3.12/dist-packages/torch/_ops.py:1243:0
  %22409 : int[] = prim::Constant[value=[1, 1280]]()
  %22477 : Double(device=cpu) = prim::Constant[value={1.19209e-07}](), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.linear.Linear::classifier
  %22478 : Long(device=cpu) = prim::Constant[value={6}](), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.linear.Linear::classifier
  %scale : Float(1, strides=[1], requires_grad=0, device=cpu), %zero_point : Char(1, strides=[1], requires_grad=0, device=cpu) = torchao::choose_qparams_affine(%input_tensor, %21943, %22409, %22464, %18935, %18935, %22477, %22478, %22464), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.linear.Linear::classifier # /usr/local/lib/python3.12/dist-packages/torch/_ops.py:1243:0
  %int_data : Char(1, 1280, strides=[1280, 1], requires_grad=0, device=cpu) = torchao::quantize_affine(%input_tensor, %22409, %scale, %zero_point, %22464, %18935, %18935), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.linear.Linear::classifier # /usr/local/lib/python3.12/dist-packages/torch/_ops.py:1243:0
  %21975 : Float(1, 1280, strides=[1280, 1], requires_grad=0, device=cpu) = torchao::dequantize_affine(%int_data, %22409, %scale, %zero_point, %22464, %18935, %18935, %22478), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.linear.Linear::classifier # /usr/local/lib/python3.12/dist-packages/torch/_ops.py:1243:0
  %22413 : Float(6, 1280, strides=[1280, 1], requires_grad=0, device=cpu) = prim::Constant[value=<Tensor>]()
  %21987 : Float(1, 6, strides=[6, 1], requires_grad=1, device=cpu) = aten::linear(%21975, %22413, %classifier.bias), scope: timm.models.mobilenetv3.MobileNetV3::/torch.nn.modules.linear.Linear::classifier # /usr/local/lib/python3.12/dist-packages/torchao/dtypes/affine_quantized_tensor_ops.py:296:0
  return (%21987)

Traceback (most recent call last):
  File "/mnt/data/qiuzhewei.qzw/train/eleme/pytorch-image-models/train_qat_torchao.py", line 1432, in <module>
    main()
  File "/mnt/data/qiuzhewei.qzw/train/eleme/pytorch-image-models/train_qat_torchao.py", line 1121, in main
    to_onnx(model, epoch=epoch)
  File "/mnt/data/qiuzhewei.qzw/train/eleme/pytorch-image-models/train_qat_torchao.py", line 1127, in to_onnx
    onnx_export(
  File "/mnt/data/qiuzhewei.qzw/train/eleme/pytorch-image-models/timm/utils/onnx.py", line 82, in onnx_export
    torch.onnx.export(
  File "/usr/local/lib/python3.12/dist-packages/torch/onnx/__init__.py", line 424, in export
    export(
  File "/usr/local/lib/python3.12/dist-packages/torch/onnx/utils.py", line 522, in export
    _export(
  File "/usr/local/lib/python3.12/dist-packages/torch/onnx/utils.py", line 1457, in _export
    graph, params_dict, torch_out = _model_to_graph(
                                    ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/onnx/utils.py", line 1084, in _model_to_graph
    graph = _optimize_graph(
            ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/onnx/utils.py", line 659, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/onnx/utils.py", line 1738, in _run_symbolic_function
    raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: ONNX export failed on an operator with unrecognized namespace {op_name}. If you are trying to export a custom operator, make sure you registered it with the right domain and version.

And the conversion function is the API from timm:

def to_onnx(model, input_size=224, epoch=0):
    # save to onnx
    
    model.to('cpu')
    onnx_export(
        model,
        f"./onnx_models/eleme/mbv4l/models_qat_int8_epoch{epoch}.onnx",
        opset=18,
        dynamic_size=False,
        aten_fallback=False,
        keep_initializers=False,
        check_forward=False,
        training=False,
        verbose=False,
        use_dynamo=False,
        input_size=(3, input_size, input_size),
        batch_size=1,
    )

Any idea how to fix that?
Thanks in advance!!

[supriyar]

cc @andrewor14

[andrewor14]

Hi @qiuzhewei, I'm not sure if we support exporting to onnx even for PTQ actually, and QAT follows the same format as PTQ. Here's a related thread: #777. @jerryzh168 do you know the latest on this?

[qiuzhewei]

My purpose is to deploy INT8 model using tritonserver, but before that, to my knowledge, I have to generate a onnx version of original model, maybe I am wrong, if there is any other way to play around?