Implement pycuda backend for inference with TensorRT engine

in additon to original polygraphy backend.
the default is polygraphy. You can set TRT_BACKEND to 'PYCUDA' to choose pycuda backend.

rtmo_gpu.py

@@ -5,6 +5,7 @@ import onnxruntime as ort
 import cv2
 from queue import Queue
 os.environ['ORT_TENSORRT_EXTRA_PLUGIN_LIB_PATHS']='libmmdeploy_tensorrt_ops.so'
+TRT_BACKEND='POLYGRAPHY'
 
 # dictionary from https://github.com/Tau-J/rtmlib/blob/4b29101d54b611048ef165277cebfffff3030074/rtmlib/visualization/skeleton/coco17.py
 coco17 = dict(name='coco17',
@@ -442,17 +443,39 @@ class RTMO_GPU(object):
                 outputs = [output.numpy() for output in io_binding.get_outputs()]
 
             else: # 'engine'
-
-                if not self.session.is_active:
-                    self.session.activate()
-
-                outputs = self.session.infer(feed_dict={'input': input}, check_inputs=False)
-                outputs = [output for output in outputs.values()]
+                if TRT_BACKEND == 'POLYGRAPHY':
+                    if not self.session.is_active:
+                        self.session.activate()
+
+                    outputs = self.session.infer(feed_dict={'input': input}, check_inputs=False)
+                    outputs = [output for output in outputs.values()]
+                else: # PYCUDA
+                    import pycuda.driver as cuda
+                    # Set the input shape dynamically
+                    input_shape = input.shape
+                    self.context.set_binding_shape(0, input_shape)
+
+                    # Ensure input_data matches the expected shape
+                    np.copyto(self.inputs[0]['host'], input.ravel())
+                    cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
+                    
+                    # Run inference
+                    self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
+                    
+                    # Transfer predictions back from the GPU
+                    for output in self.outputs:
+                        cuda.memcpy_dtoh_async(output['host'], output['device'], self.stream)
+                    
+                    # Synchronize the stream
+                    self.stream.synchronize()
+                    
+                    # Return only the output values (in their original shapes)
+                    outputs = [out['host'].reshape(out['shape']) for out in self.outputs]
 
             return outputs
 
     def __exit__(self):
-        if self.model_format == 'engine':
+        if self.model_format == 'engine' and TRT_BACKEND == 'POLYGRAPHY':
             if self.session.is_active:
                 self.session.deactivate()
 
@@ -471,7 +494,11 @@ class RTMO_GPU(object):
                  mean: tuple = None,
                  std: tuple = None,
                  device: str = 'cuda',
-                 is_yolo_nas_pose = False):
+                 is_yolo_nas_pose = False, 
+                 batch_size = 1,
+                 plugin_path = 'libmmdeploy_tensorrt_ops.so'):
+
+        self.batch_size = batch_size
 
         if not os.path.exists(model):
             # If the file does not exist, raise FileNotFoundError
@@ -499,10 +526,62 @@ class RTMO_GPU(object):
                                                 providers=providers[device])
 
         else: # 'engine'
-            from polygraphy.backend.common import BytesFromPath
-            from polygraphy.backend.trt import EngineFromBytes, TrtRunner
-            engine = EngineFromBytes(BytesFromPath(model))
-            self.session = TrtRunner(engine)
+            if TRT_BACKEND == 'POLYGRAPHY':
+                from polygraphy.backend.common import BytesFromPath
+                from polygraphy.backend.trt import EngineFromBytes, TrtRunner
+                engine = EngineFromBytes(BytesFromPath(model))
+                self.session = TrtRunner(engine)
+            else: # PYCUDA
+                import tensorrt as trt
+                import ctypes
+                import pycuda.autoinit
+                import pycuda.driver as cuda
+                self.TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+                self.trt_model_path = model
+                self.plugin_path = plugin_path
+
+                # Load the custom plugin library
+                ctypes.CDLL(self.plugin_path)
+
+                # Load the TensorRT engine
+                with open(self.trt_model_path, 'rb') as f:
+                    engine_data = f.read()
+                
+                self.runtime = trt.Runtime(self.TRT_LOGGER)
+                self.engine = self.runtime.deserialize_cuda_engine(engine_data)
+
+                if self.engine is None:
+                    raise RuntimeError("Failed to load the engine.")
+
+                self.context = self.engine.create_execution_context()
+
+                self.inputs = []
+                self.outputs = []
+                self.bindings = []
+                self.stream = cuda.Stream()
+
+                # Allocate memory for inputs and outputs
+                for binding in self.engine:
+                    binding_index = self.engine.get_binding_index(binding)
+                    shape = self.engine.get_binding_shape(binding_index)
+                    if shape[0] == -1:
+                        # Handle dynamic batch size by setting max_batch_size
+                        shape[0] = self.batch_size
+                    size = trt.volume(shape)
+                    dtype = trt.nptype(self.engine.get_binding_dtype(binding))
+                    
+                    # Allocate host and device buffers
+                    host_mem = cuda.pagelocked_empty(size, dtype)
+                    device_mem = cuda.mem_alloc(host_mem.nbytes)
+                    
+                    # Append the device buffer to device bindings.
+                    self.bindings.append(int(device_mem))
+                    
+                    # Append to the appropriate list.
+                    if self.engine.binding_is_input(binding):
+                        self.inputs.append({'host': host_mem, 'device': device_mem, 'shape': shape})
+                    else:
+                        self.outputs.append({'host': host_mem, 'device': device_mem, 'shape': shape})
 
         self.model_input_size = self.input_shape[2:4] # B, C, H, W,
         self.mean = mean
@@ -510,6 +589,9 @@ class RTMO_GPU(object):
         self.device = device
         self.is_yolo_nas_pose = is_yolo_nas_pose
 
+        print(f'[I] Detected \'{self.model_format.upper()}\' model', end='')
+        print(f', \'{TRT_BACKEND.upper()}\' backend is chosen for inference' if self.model_format == 'engine' else '')
+        
 class RTMO_GPU_Batch(RTMO_GPU):
     def preprocess_batch(self, imgs: List[np.ndarray]) -> Tuple[np.ndarray, List[float]]:
         """Process a batch of images for RTMPose model inference.
@@ -571,12 +653,34 @@ class RTMO_GPU_Batch(RTMO_GPU):
             outputs = [output.numpy() for output in io_binding.get_outputs()]
 
         else: # 'engine'
+            if TRT_BACKEND == 'POLYGRAPHY':
+                if not self.session.is_active:
+                    self.session.activate()
 
-            if not self.session.is_active:
-                self.session.activate()
-
-            outputs = self.session.infer(feed_dict={'input': input}, check_inputs=False)
-            outputs = [output for output in outputs.values()]
+                outputs = self.session.infer(feed_dict={'input': input}, check_inputs=False)
+                outputs = [output for output in outputs.values()]
+            else: # PYCUDA
+                import pycuda.driver as cuda
+                # Set the input shape dynamically
+                input_shape = input.shape
+                self.context.set_binding_shape(0, input_shape)
+
+                # Ensure input_data matches the expected shape
+                np.copyto(self.inputs[0]['host'], input.ravel())
+                cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
+                
+                # Run inference
+                self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
+                
+                # Transfer predictions back from the GPU
+                for output in self.outputs:
+                    cuda.memcpy_dtoh_async(output['host'], output['device'], self.stream)
+                
+                # Synchronize the stream
+                self.stream.synchronize()
+                
+                # Return only the output values (in their original shapes)
+                outputs = [out['host'].reshape(out['shape']) for out in self.outputs]
 
         return outputs
 
@@ -651,14 +755,16 @@ class RTMO_GPU_Batch(RTMO_GPU):
                  std: tuple = None,
                  device: str = 'cuda',
                  is_yolo_nas_pose = False,
+                 plugin_path = 'libmmdeploy_tensorrt_ops.so',
                  batch_size: int = 1):
         super().__init__(model, 
                          mean, 
                          std, 
                          device, 
-                         is_yolo_nas_pose)
+                         is_yolo_nas_pose, 
+                         batch_size,
+                         plugin_path)
         
-        self.batch_size = batch_size
         self.in_queues = dict()
         self.out_queues = dict()
         self.buffers = dict()