Error while creating tensorrt engine for SAM Vision Transformer using timm library

[deo-abhijit]

Error Log:

2025-01-13 17:19:30.473927780 [W:onnxruntime:Default, onnxruntime_pybind_state.cc:965 CreateExecutionProviderInstance] Failed to create CUDAExecutionProvider. Require cuDNN 9.* and CUDA 12.*. Please install all dependencies as mentioned in the GPU requirements page (https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements), make sure they're in the PATH, and that your GPU is supported.


ONNX model tested and results look good!
[01/13/2025-17:19:40] [TRT] [W] UNSUPPORTED_STATE: Skipping tactic 0 due to insufficient memory on requested size of 1803374592 detected for tactic 0x0000000000000000.
[01/13/2025-17:19:40] [TRT] [E] IBuilder::buildSerializedNetwork: Error Code 10: Internal Error (Could not find any implementation for node {ForeignNode[/blocks/blocks.10/attn/Gather_2_output_0[Constant].../blocks/blocks.11/Reshape_1 + /Transpose]}.)
Traceback (most recent call last):
  File "/home/mzcar/accelerated/test_script.py", line 74, in <module>
    engine = backend.prepare(onnx_model, device='CUDA:0')
  File "/home/mzcar/miniconda3/envs/mzinferx/lib/python3.10/site-packages/onnx_tensorrt/backend.py", line 236, in prepare
    return TensorRTBackendRep(model, device, **kwargs)
  File "/home/mzcar/miniconda3/envs/mzinferx/lib/python3.10/site-packages/onnx_tensorrt/backend.py", line 92, in __init__
    self._build_engine()
  File "/home/mzcar/miniconda3/envs/mzinferx/lib/python3.10/site-packages/onnx_tensorrt/backend.py", line 132, in _build_engine
    raise RuntimeError("Failed to build TensorRT engine from network")
RuntimeError: Failed to build TensorRT engine from network

TensorRT Version:

pip show onnx-tensorrt
Name: onnx_tensorrt
Version: 10.7.0
Summary: ONNX-TensorRT - TensorRT backend for running ONNX models
Home-page: https://github.com/onnx/onnx-tensorrt
Author: NVIDIA
Author-email: [email protected]
License: 
Location: /home/mzcar/miniconda3/envs/mzinferx/lib/python3.10/site-packages
Requires: numpy, onnx, pycuda
Required-by: 

Code to reprodue:

import timm 
import torch, pathlib,onnx,onnxruntime 
import numpy as np
import onnx_tensorrt.backend as backend

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

class OnnxExporter:
    """
    
    """
    def __init__(self,model,**kwargs):
        super().__init__()
        self.model_name=type(model).__name__
        self.model = model
        self.input_image_size = (1,3,1024,1024) 
        self.onnx_model_path = pathlib.Path(f'~/.cache/onnx_models/{self.model_name}/{self.model_name}.onnx').expanduser()
        self.onnx_model_path.parent.mkdir(exist_ok=True, parents=True)
        self.input_tensor = torch.randn(self.input_image_size)


    def convert_to_onnx(self):
        self._onnx_export() 
        if not self._check_if_export_okay():
            raise RuntimeError("Could not export onnx model")
        
        torch_out = self.model(self.input_tensor)
        ort_session = onnxruntime.InferenceSession(self.onnx_model_path, providers=['CPUExecutionProvider',"CUDAExecutionProvider"])
        ort_inputs = {'input':to_numpy(self.input_tensor)}
        ort_outs = ort_session.run(None,ort_inputs)
        
        np.testing.assert_allclose(to_numpy(torch_out).squeeze(), ort_outs[0].squeeze(), rtol=1e-3, atol=1e-3)
        print('ONNX model tested and results look good!')
        return self.onnx_model_path
        
    def _onnx_export(self):
        print(self.onnx_model_path)
        if self.onnx_model_path.is_file():
            print("path already exists")
            return         
        torch.onnx.export(
            model=self.model,
            args=self.input_tensor,
            f=pathlib.Path(f'~/.cache/onnx_models/{self.model_name}/{self.model_name}.onnx').expanduser(),
            export_params=True,
            opset_version=17,
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
        )

    def _check_if_export_okay(self,):
        try:
            onnx_model= onnx.load(self.onnx_model_path)
            status=onnx.checker.check_model(onnx_model)
            return True 
        except:
            return False


model = timm.create_model(f'samvit_base_patch16.sa1b',pretrained=False)

onnx_exporter = OnnxExporter(model=model)
onnx_path = onnx_exporter.convert_to_onnx()

onnx_model = onnx.load(onnx_path)
engine = backend.prepare(onnx_model, device='CUDA:0')

I am able to export the model to ONNX, but converting to TensorRT engine is failing. Could anyone help me out here?

Thanks!