Skip to content

[Bug] [FRONTEND][ONNX] failed to import TopK operator when K is a constant tensor: ValueError: TopK k must be a constant #17856

Closed
@coffezhou

Description

@coffezhou

Expected behavior

The onnx frontend should import the model correctly.

Actual behavior

When importing the following onnx model, tvm crashes as follows:

Image

Traceback (most recent call last):
  File "/home/carla/Documents/test/test.py", line 189, in <module>
    main()
  File "/home/carla/Documents/test/test.py", line 178, in main
    check_correctness(onnx_model)
  File "/home/carla/Documents/test/test.py", line 104, in check_correctness
    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3690, in from_onnx
    return g.from_onnx(graph, opset)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3321, in from_onnx
    self._construct_nodes(graph)
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3496, in _construct_nodes
    op = self._convert_operator(op_name, inputs, attr, self.opset)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3596, in _convert_operator
    sym = op_function(self.bb, inputs, attrs, [self._nodes, self._params])
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 2706, in _impl_v11
    raise ValueError("TopK k must be a constant")
ValueError: TopK k must be a constant
[17:35:01] /home/carla/Documents/tvm/src/relax/ir/block_builder.cc:64: Warning: BlockBuilder destroyed with remaining blocks!

The attribute K of TopK should be a constant. In this model, K is a constant tensor, which meets the constraint of TopK. Thus, this may be a bug of the onnx frontend in TVM.

Environment

OS: Ubuntu 20.04
TVM: 0.21.dev0(c00f52a)

Steps to reproduce

This bug can be reproduced by the following code with the model in the attachment. As shown in the code, the model can be executed by onnxruntime.

import sys

import numpy as np
import onnx
import onnxruntime

import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

import pickle

            
def main():
    onnx_model = onnx.load("a157.onnx")
    
    with open("inputs.pkl", "rb") as fp:
        inputs = pickle.load(fp)
    
    try:
        ort_session = onnxruntime.InferenceSession(
            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        ort_output = ort_session.run([], inputs)
    except Exception as e:
        print(e)
        sys.exit(1)
        
    # Convert the onnx model into relax through the onnx importer.
    tvm_model = from_onnx(onnx_model, keep_params_in_input=True)
   
if __name__ == "__main__":
    
    main()

testcase.zip

Triage

  • needs-triage

cc @KJlaccHoeUM9l

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions