EfficientViT-SAM TensorRT export error due to invalid node (temp solution by monkey patching original sam repo)

[xuanlinli17]

When converting the EfficientViT-SAM model from torch to ONNX to TensorRT with latest torch version (2.2), I encounter the following errors while converting from ONNX to TensorRT:

[03/08/2024-23:29:49] [E] [TRT] ModelImporter.cpp:771: While parsing node number 108 [Tile -> "/Tile_output_0"]:
[03/08/2024-23:29:49] [E] [TRT] ModelImporter.cpp:772: --- Begin node ---
[03/08/2024-23:29:49] [E] [TRT] ModelImporter.cpp:773: input: "/Unsqueeze_3_output_0"
input: "/Reshape_2_output_0"
output: "/Tile_output_0"
name: "/Tile"
op_type: "Tile"

Upon further debugging, I find that the torch.repeat_interleave operations in the original SAM repo (https://github.com/facebookresearch/segment-anything/blob/6fdee8f2727f4506cfbbe553e23b895e27956588/segment_anything/modeling/mask_decoder.py#L112) is causing the issue.

The temporary solution is to monkey-patch the functions in original SAM that utilize torch.repeat_interleave with some equivalent code using torch.tile see e.g., https://github.com/xuanlinli17/efficientvit/commit/322ce5ea4cdbd61af2efcf3710f052b0d0eeed54