facebookresearch / facebookresearch/sam3
The TensorRT result of ViT PatchEmbed module is incorrect
- Dominant language
- Python
- Stars
- 11.7k
- Forks
- 1.8k
- PR merge metrics
- No merged PRs in 30d
Description
When I convert ViT's PatchEmbed module from onnx to TensorRT, their result is very different. Here is my code:
```python
import os
import torch
import onnxruntime
import tensorrt as trt
from PIL import Image
from torchvision.transforms import v2
from sam3.model_builder import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor
class VITSub(torch.nn.Module):
def __init__(self, model):
super(VITSub, self).__init__()
self.model = model.backbone.vision_backbone.trunk
def forward(self, x):
x = self.model.patch_embed(x)
return x
if __name__ == "__main__":
onnx_path = 'vit_sub.onnx'
engine_path = 'vit_sub.engine'
model = build_sam3_image_model(load_from_HF=False, checkpoint_path=./sam3.pt", device='cuda')
model.eval()
image = Image.open("assets/images/test_image.jpg")
image = v2.functional.to_image(image).to('cuda')
processor = Sam3Processor(model)
image = processor.transform(image).unsqueeze(0)
torch.onnx.export(
model.backbone.vision_backbone.trunk.patch_embed,
image,
onnx_path,
input_names=['image'],
# dynamic_axes={name: {0: "batch"} for name in (input_names + output_names)},
output_names=['output'])
sess = onnxruntime.InferenceSession(onnx_path, providers=["CUDAExecutionProvider"])
r_out = sess.run( ['output'], {'image': image.cpu().numpy()})
cmd = f"trtexec --onnx={onnx_path} --saveEngine={engine_path}"
os.system(cmd)
with open(engine_path, "rb") as f, trt.Runtime(trt.Logger()) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
output = torch.empty(1, 72, 72, 1024, dtype=torch.float32, device='cuda')
bindings = [image.data_ptr(), output.data_ptr()]
context.execute_v2(bindings=bindings)
max_diff =(torch.from_numpy(r_out[0]).to(output) - output).abs().max()
print('max_diff: ', max_diff) # my result is 'max_diff: max_diff: tensor(7.6664, device='cuda:0')
```
my environment is that:
```shell
torch 2.6.0
tensorrt 10.14.1.48
onnx 1.17.0
onnxruntime-gpu 1.19.0
Driver Version: 570.133.07
CUDA Version: 12.8
NVIDIA GeForce RTX 3090
```
Contributor guide
Research direction
Reproduce the conversion using the provided VITSub example, build_sam3_image_model entry point, and assets/images/test_image.jpg. Compare the ONNX Runtime and TensorRT outputs for trunk.patch_embed, then inspect the TensorRT conversion behavior for this module. Done means the output discrepancy is explained and a verified correction or clear compatibility limitation is documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 32/100