From 28069f63f991e22a3b9b36184629e17e9e7303bc Mon Sep 17 00:00:00 2001 From: Alex Bezdieniezhnykh Date: Wed, 30 Apr 2025 23:47:46 +0300 Subject: [PATCH] Reapply "import Tensorrt not in compile time in order to dynamically load tensorrt only if nvidia gpu is present" This reverts commit cf01e5d95212062e35d78982a73f0fc9b11444ea. --- Azaion.Inference/inference.pyx | 2 +- Azaion.Inference/onnx_engine.pyx | 4 ++-- Azaion.Inference/tensorrt_engine.pxd | 14 +++----------- Azaion.Inference/tensorrt_engine.pyx | 12 ++++++------ 4 files changed, 12 insertions(+), 20 deletions(-) diff --git a/Azaion.Inference/inference.pyx b/Azaion.Inference/inference.pyx index d505e60..48925de 100644 --- a/Azaion.Inference/inference.pyx +++ b/Azaion.Inference/inference.pyx @@ -16,7 +16,7 @@ from hardware_service cimport HardwareService from security cimport Security if HardwareService.has_nvidia_gpu(): - from tensorrt_engine cimport TensorRTEngine + from tensorrt_engine import TensorRTEngine else: from onnx_engine import OnnxEngine diff --git a/Azaion.Inference/onnx_engine.pyx b/Azaion.Inference/onnx_engine.pyx index ed8cdeb..de68d73 100644 --- a/Azaion.Inference/onnx_engine.pyx +++ b/Azaion.Inference/onnx_engine.pyx @@ -14,11 +14,11 @@ cdef class OnnxEngine(InferenceEngine): model_meta = self.session.get_modelmeta() print("Metadata:", model_meta.custom_metadata_map) - cdef tuple get_input_shape(self): + cpdef tuple get_input_shape(self): shape = self.input_shape return shape[2], shape[3] - cdef int get_batch_size(self): + cpdef int get_batch_size(self): return self.batch_size cpdef run(self, input_data): diff --git a/Azaion.Inference/tensorrt_engine.pxd b/Azaion.Inference/tensorrt_engine.pxd index 90d8b79..6fc31bd 100644 --- a/Azaion.Inference/tensorrt_engine.pxd +++ b/Azaion.Inference/tensorrt_engine.pxd @@ -16,17 +16,9 @@ cdef class TensorRTEngine(InferenceEngine): cdef object stream - @staticmethod - cdef get_gpu_memory_bytes(int device_id) - @staticmethod - cdef get_engine_filename(int device_id) + cpdef tuple get_input_shape(self) - @staticmethod - cdef convert_from_onnx(bytes onnx_model) + cpdef int get_batch_size(self) - cdef tuple get_input_shape(self) - - cdef int get_batch_size(self) - - cdef run(self, input_data) + cpdef run(self, input_data) diff --git a/Azaion.Inference/tensorrt_engine.pyx b/Azaion.Inference/tensorrt_engine.pyx index c7cf1a5..d8c4189 100644 --- a/Azaion.Inference/tensorrt_engine.pyx +++ b/Azaion.Inference/tensorrt_engine.pyx @@ -56,7 +56,7 @@ cdef class TensorRTEngine(InferenceEngine): raise RuntimeError(f"Failed to initialize TensorRT engine: {str(e)}") @staticmethod - cdef get_gpu_memory_bytes(int device_id): + def get_gpu_memory_bytes(int device_id): total_memory = None try: pynvml.nvmlInit() @@ -73,7 +73,7 @@ cdef class TensorRTEngine(InferenceEngine): return 2 * 1024 * 1024 * 1024 if total_memory is None else total_memory # default 2 Gb @staticmethod - cdef get_engine_filename(int device_id): + def get_engine_filename(int device_id): try: device = cuda.Device(device_id) sm_count = device.multiprocessor_count @@ -83,7 +83,7 @@ cdef class TensorRTEngine(InferenceEngine): return None @staticmethod - cdef convert_from_onnx(bytes onnx_model): + def convert_from_onnx(bytes onnx_model): workspace_bytes = int(TensorRTEngine.get_gpu_memory_bytes(0) * 0.9) explicit_batch_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) @@ -112,13 +112,13 @@ cdef class TensorRTEngine(InferenceEngine): constants.log('conversion done!') return bytes(plan) - cdef tuple get_input_shape(self): + cpdef tuple get_input_shape(self): return self.input_shape[2], self.input_shape[3] - cdef int get_batch_size(self): + cpdef int get_batch_size(self): return self.batch_size - cdef run(self, input_data): + cpdef run(self, input_data): try: cuda.memcpy_htod_async(self.d_input, input_data, self.stream) self.context.set_tensor_address(self.input_name, int(self.d_input)) # input buffer