grid_sampler layer and quantization with onnxruntime doesn't work

[avostryakov]

Describe the Issue
I converted a model with grid_sampler layer from pytorch 1.9 / 1.10 to onnx format. It works and I can open it in onnxruntime with help of mmcv. But when I make quantization of the model (succesfully) I can't run it in onnxruntime with the following error:

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from var/pretrained/kp2d-v4-quantized.onnx failed:This is an invalid model. Error: Duplicate definition of name (input).
where input - is input variable name

Also, during quantization I see the following warnings:

Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator MaxPool. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator MaxPool. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator ReorderOutput. No schema registered for this operator.
Warning: Unsupported operator MaxPool. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator ReorderOutput. No schema registered for this operator.
Warning: Unsupported operator ReorderInput. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator ReorderOutput. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator Conv. No schema registered for this operator.
Warning: Unsupported operator ReorderOutput. No schema registered for this operator.
Warning: Unsupported operator ReorderOutput. No schema registered for this operator.

Reproduction

Quantization code:

from onnxruntime.quantization import quantize_dynamic, QuantType

model_name = onnx_model_path.as_posix()
output_model_name = f"{model_name[:-5]}-quantized.onnx"
quantize_dynamic(
    model_input=model_name,
    model_output=output_model_name,
    per_channel=True,
    activation_type=QuantType.QUInt8,
    weight_type=QuantType.QUInt8,
    optimize_model=False,
)

Environment

Ubuntu 20.04
pytorch 1.9.1/1.10
onnx 1.8.1
onnxruntime 1.8.1
mmcv-full 1.3.15 + compiled extensions _ext*.so for Linux

Question

Does it mean that quantization of models with grid_sampler inside don't support or I do something wrong?

[grimoire]

Hi, the schema of an ONNX node is usually used to support shape inference.
We have added the symbol of grid sample and implementation of the ops, but we haven't added the schema for it(of cause, we are trying).
I guess you have exported the ONNX with a dynamic shape, right? If you do so, try export it with a static shape and try again.
Or you can add the shape info (model.graph.value_info) before quantization and remove then afterward.

[avostryakov]

Hi, Grimoire!

Hi, the schema of an ONNX node is usually used to support shape inference.
We have added the symbol of grid sample and implementation of the ops, but we haven't added the schema for it(of cause, we are trying).
I guess you have exported the ONNX with a dynamic shape, right? If you do so, try export it with a static shape and try again.

Let me clarify what do you mean by static shape. I used standard pytorch.onnx.export:

    input_names = ['input']
    input_dynamic_axes = {'input': {0: 'batch', 1: 'channel', 2: 'height', 3: 'width'}}

    output_names = ['scores_raw', 'coords_raw', 'descr_raw']
    output_dynamic_axes = {'scores_raw': {0: 'batch', 1: 'score', 2: 'height', 3: 'width'},
                                              'coords_raw': {0: 'batch', 1: 'score', 2: 'height', 3: 'width'},
                                              'descr_raw': {0: 'batch', 1: 'score', 2: 'height', 3: 'width'}}
    dynamic_axes = dict(input_dynamic_axes, **output_dynamic_axes)
    
    export(
            model,
            (input_var,),
            f=output.as_posix(),
            input_names=input_names,
            output_names=output_names,
            dynamic_axes=dynamic_axes,
            do_constant_folding=True,
            use_external_data_format=use_external_format,
            enable_onnx_checker=True,
            opset_version=11,
        )

Is it static or dynamic in my case?

Or you can add the shape info (model.graph.value_info) before quantization and remove then afterward.

Do you mean to load a model with onnx:
onnx_model = onnx.load(onnx_model_path.as_posix())

and add this shape info?

[grimoire]

Once you have given the dynamic_axes, the given axes will be dynamic.
If the graph is static, there will be information of tensor for each layer in onnx_model.graph.value_info.

[avostryakov]

Once you have given the dynamic_axes, the given axes will be dynamic.
If the graph is static, there will be information of tensor for each layer in onnx_model.graph.value_info.

I tried to export without dynamic_axes parameter. It doesn't work too when I try to execute inference after quantization. Before quantization inference of onnx model works. Maybe it starts to work when grid_sampler will be added to onnx and onnxruntime library

[grimoire]

Hope it will be added to ONNX in the future.

Err, another workaround is adding the schema by yourself. ONNX defines the schema in schema.h and register it in operator_set.h. I have not done it before.