mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 15:06:29 +00:00
split to 2 files: tensorrt_engine and onnx engine
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user