SAM-vit_b-prompt model build engine failed of TensorRT 8.6.12 when using dynamic shape of num_point

[haibozhang123]

Description

Environment

TensorRT Version: 8.6.12

NVIDIA GPU:

NVIDIA Driver Version:

CUDA Version:

CUDNN Version:

Operating System:

Python Version (if applicable): 3.9

Tensorflow Version (if applicable):

PyTorch Version (if applicable):

Baremetal or Container (if so, version):

Relevant Files

Model link:

Steps To Reproduce

1. git clone https://github.com/facebookresearch/segment-anything

2. export prompt_mask

`def export_prompt_masks_model(model_type: str, checkpoint: str, opset: 17):
print("Loading model...")
sam = sam_model_registry“vit_b”

onnx_model = SamOnnxModel(
    model=sam,
    return_single_mask=False,
    # use_stability_score=False,
    # return_extra_metrics=False,
)

onnx_model_path = os.path.join("weights", "sam_" + model_type+"_"+"prompt_mask_format_dynamic6.onnx")

dynamic_axes = {
    "point_coords": {1: "num_points"},
    "point_labels": {1: "num_points"},
}

embed_dim = sam.prompt_encoder.embed_dim
embed_size = sam.prompt_encoder.image_embedding_size
mask_input_size = [4 * x for x in embed_size]
dummy_inputs = {
    "image_embeddings": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),
    "point_coords": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),
    "point_labels": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),
    "mask_input": torch.randn(1, 1, *mask_input_size, dtype=torch.float),
    "has_mask_input": torch.tensor([1], dtype=torch.float),
    "orig_im_size": torch.tensor([1500, 2250], dtype=torch.float),
}

_ = onnx_model(**dummy_inputs)

output_names = ["masks", "iou_predictions", "low_res_masks"]

with warnings.catch_warnings():
    warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
    warnings.filterwarnings("ignore", category=UserWarning)
    with open(onnx_model_path, "wb") as f:
        print(f"Exporting onnx model to {onnx_model_path}...")
        torch.onnx.export(
            onnx_model,
            tuple(dummy_inputs.values()),
            f,
            export_params=True,
            verbose=False,
            opset_version=opset,
            do_constant_folding=True,
            input_names=list(dummy_inputs.keys()),
            output_names=output_names,
            dynamic_axes=dynamic_axes,
        )`

3. trtexec import trt engine
   trtexec --onnx=sam_vit_b_dy.onnx --saveEngine=sam_vit_b_prompt_mask_format_dynamic.engine --minShapes=point_coords:1x1x2,point_labels:1x1 --optShapes=point_coords:1x1x2,point_labels:1x1 --maxShapes=point_coords:1x10x2,point_labels:1x10

4. found
   

Commands or scripts:

Have you tried the latest release?:

Can this model run on other frameworks? For example run ONNX model with ONNXRuntime (polygraphy run <model.onnx> --onnxrt):

[zerollzeng]

As it said: OneHot: Onehot layer cannot be used to compute a shape tensor.

It's a known limitation. Could you please try constant folding?

[haibozhang123]

I using constant_folding in torch-onnx export in front code ,but ...

`

 with open(onnx_model_path, "wb") as f:

    print(f"Exporting onnx model to {onnx_model_path}...")
    torch.onnx.export(
        onnx_model,
        tuple(dummy_inputs.values()),
        f,
        export_params=True,
        verbose=False,
        opset_version=opset,
        do_constant_folding=True,
        input_names=list(dummy_inputs.keys()),
        output_names=output_names,
        dynamic_axes=dynamic_axes,
    )

`

[zerollzeng]

Okay, looks like this need to fixed in Torch source code...

[zerollzeng]

Or it just simply a TRT limitation.

[ttyio]

closing since no activity for more than 3 weeks, pls reopen if you still have question, thanks!

[cxmt-ai-tc]

I using constant_folding in torch-onnx export in front code ,but ...

`

 with open(onnx_model_path, "wb") as f:

    print(f"Exporting onnx model to {onnx_model_path}...")
    torch.onnx.export(
        onnx_model,
        tuple(dummy_inputs.values()),
        f,
        export_params=True,
        verbose=False,
        opset_version=opset,
        do_constant_folding=True,
        input_names=list(dummy_inputs.keys()),
        output_names=output_names,
        dynamic_axes=dynamic_axes,
    )

`

does it work？ do_constant_folding=True,