Skip to content

TPU-MLIR model_transform fails on Mask R-CNN ONNX export — KeyError on tensor inside If subgraph #270

Description

@Ryan-py

Describe the bug

model_transform fails when converting a Mask R-CNN ONNX model. The converter raises a KeyError while parsing an If op subgraph inside the RPN, because the subgraph references a tensor (/rpn/Gather_78_output_0) defined in the outer graph scope.

Steps to reproduce

1. export a torchvision Mask R-CNN to ONNX
2. Run `model_transform`
`model_transform \
     --model_name maskrcnn \
     --model_def maskrcnn.onnx \
     --input_shapes "[[3,648,1152]]" \
     --mean "0.0,0.0,0.0" \
     --scale "1.0,1.0,1.0" \
     --pixel_format rgb \
     --test_input ../image/chicken.jpg \
     --test_result maskrcnn_top_outputs.npz \
     --mlir chicken_maskrcnn.mlir`

Expected behavior

model_transform should resolve outer-graph tensor references when parsing If subgraphs, the same way ONNX Runtime does. Per the ONNX spec, subgraph nodes are allowed to reference values from any enclosing scope.

Error logs / stack trace

`WARNING: onnxsim opt failed.
WARNING: onnx model check failed
WARNING: ConstantFolding failed.
Traceback (most recent call last):
  File ".../model_transform.py", line 292, in <module>
    tool.model_transform(args.mlir, args.add_postprocess, args.patterns_count)
  File ".../OnnxConverter.py", line 2711, in convert_if_op
    self.parse_subgraph(new_op.owner, region_idx, attr.g)
  File ".../OnnxConverter.py", line 2679, in parse_subgraph
    self.onnxop_factory.get(n.op_type, lambda x: NoneAndRaise(x))(n)
  File ".../OnnxConverter.py", line 673, in convert_cast_op
    op = self.getOperand(onnx_node.inputs[0])
  File ".../BaseConverter.py", line 54, in getOperand
    raise KeyError("operand {} not found".format(name))
KeyError: 'operand /rpn/Gather_78_output_0 not found'`

TPU-MLIR version

Name: tpu_mlir Version: 1.7

Target chip

cv18xx

Source framework / model

Pytorch

OS / Docker image

sophgo/tpuc_dev:v3.1

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions