split to 2 files: tensorrt_engine and onnx engine

This commit is contained in:
Alex Bezdieniezhnykh
2025-04-30 19:33:59 +03:00
parent 4bca61d859
commit a1ee077e0a
12 changed files with 227 additions and 206 deletions
+1 -27
View File
@@ -6,30 +6,4 @@ cdef class InferenceEngine:
cdef public int batch_size
cdef tuple get_input_shape(self)
cdef int get_batch_size(self)
cpdef run(self, input_data)
cdef class OnnxEngine(InferenceEngine):
cdef object session
cdef list model_inputs
cdef str input_name
cdef object input_shape
cdef class TensorRTEngine(InferenceEngine):
cdef object stream
cdef object context
cdef str input_name
cdef str output_name
cdef object d_input
cdef object d_output
cdef object input_shape
cdef object output_shape
cdef object h_output
@staticmethod
cdef bytes convert_from_onnx(bytes onnx_model)
@staticmethod
cdef unsigned long long get_gpu_memory_bytes(device_id=?)
@staticmethod
cdef str get_engine_filename(device_id=?)
cdef run(self, input_data)