facebookresearch / facebookresearch/sam3

The TensorRT result of ViT PatchEmbed module is incorrect

Open
#388 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.