mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 11:46:30 +00:00
clean postbuild script
clean warnings
This commit is contained in:
@@ -7,7 +7,7 @@ import tensorrt as trt
|
||||
import pycuda.driver as cuda
|
||||
import pycuda.autoinit # required for automatically initialize CUDA, do not remove.
|
||||
import pynvml
|
||||
|
||||
cimport constants
|
||||
|
||||
cdef class InferenceEngine:
|
||||
def __init__(self, model_bytes: bytes, batch_size: int = 1, **kwargs):
|
||||
@@ -51,7 +51,6 @@ cdef class OnnxEngine(InferenceEngine):
|
||||
cdef class TensorRTEngine(InferenceEngine):
|
||||
def __init__(self, model_bytes: bytes, batch_size: int = 4, **kwargs):
|
||||
super().__init__(model_bytes, batch_size)
|
||||
print('Enter init TensorRT')
|
||||
try:
|
||||
logger = trt.Logger(trt.Logger.WARNING)
|
||||
|
||||
@@ -142,7 +141,7 @@ cdef class TensorRTEngine(InferenceEngine):
|
||||
return None
|
||||
|
||||
if builder.platform_has_fast_fp16:
|
||||
print('Converting to supported fp16')
|
||||
constants.log('Converting to supported fp16')
|
||||
config.set_flag(trt.BuilderFlag.FP16)
|
||||
else:
|
||||
print('Converting to supported fp32. (fp16 is not supported)')
|
||||
@@ -151,7 +150,7 @@ cdef class TensorRTEngine(InferenceEngine):
|
||||
if plan is None:
|
||||
print('Conversion failed.')
|
||||
return None
|
||||
print('conversion done!')
|
||||
constants.log('conversion done!')
|
||||
return bytes(plan)
|
||||
|
||||
cdef tuple get_input_shape(self):
|
||||
|
||||
Reference in New Issue
Block a user