119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
import pycuda.autoinit
|
||
import pycuda.driver as cuda
|
||
import tensorrt as trt
|
||
import numpy as np
|
||
|
||
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
|
||
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
|
||
def GiB(val):
|
||
return val * 1 << 30
|
||
|
||
def ONNX_to_TRT(onnx_model_path=None,trt_engine_path=None,fp16_mode=False):
|
||
"""
|
||
仅适用TensorRT V8版本
|
||
生成cudaEngine，并保存引擎文件(仅支持固定输入尺度)
|
||
|
||
fp16_mode: True则fp16预测
|
||
onnx_model_path: 将加载的onnx权重路径
|
||
trt_engine_path: trt引擎文件保存路径
|
||
"""
|
||
builder = trt.Builder(TRT_LOGGER)
|
||
network = builder.create_network(EXPLICIT_BATCH)
|
||
parser = trt.OnnxParser(network, TRT_LOGGER)
|
||
|
||
config = builder.create_builder_config()
|
||
config.max_workspace_size=GiB(1)
|
||
if fp16_mode:
|
||
config.set_flag(trt.BuilderFlag.FP16)
|
||
with open(onnx_model_path, 'rb') as model:
|
||
assert parser.parse(model.read())
|
||
serialized_engine=builder.build_serialized_network(network, config)
|
||
|
||
|
||
with open(trt_engine_path, 'wb') as f:
|
||
f.write(serialized_engine) # 序列化
|
||
|
||
print('TensorRT file in ' + trt_engine_path)
|
||
print('============ONNX->TensorRT SUCCESS============')
|
||
|
||
class TrtModel():
|
||
'''
|
||
TensorRT infer
|
||
'''
|
||
def __init__(self,trt_path):
|
||
self.ctx=cuda.Device(0).make_context()
|
||
stream = cuda.Stream()
|
||
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
|
||
runtime = trt.Runtime(TRT_LOGGER)
|
||
|
||
# Deserialize the engine from file
|
||
with open(trt_path, "rb") as f:
|
||
engine = runtime.deserialize_cuda_engine(f.read())
|
||
context = engine.create_execution_context()
|
||
|
||
host_inputs = []
|
||
cuda_inputs = []
|
||
host_outputs = []
|
||
cuda_outputs = []
|
||
bindings = []
|
||
|
||
for binding in engine:
|
||
print('bingding:', binding, engine.get_binding_shape(binding))
|
||
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
|
||
dtype = trt.nptype(engine.get_binding_dtype(binding))
|
||
# Allocate host and device buffers
|
||
host_mem = cuda.pagelocked_empty(size, dtype)
|
||
cuda_mem = cuda.mem_alloc(host_mem.nbytes)
|
||
# Append the device buffer to device bindings.
|
||
bindings.append(int(cuda_mem))
|
||
# Append to the appropriate list.
|
||
if engine.binding_is_input(binding):
|
||
self.input_w = engine.get_binding_shape(binding)[-1]
|
||
self.input_h = engine.get_binding_shape(binding)[-2]
|
||
host_inputs.append(host_mem)
|
||
cuda_inputs.append(cuda_mem)
|
||
else:
|
||
host_outputs.append(host_mem)
|
||
cuda_outputs.append(cuda_mem)
|
||
|
||
# Store
|
||
self.stream = stream
|
||
self.context = context
|
||
self.engine = engine
|
||
self.host_inputs = host_inputs
|
||
self.cuda_inputs = cuda_inputs
|
||
self.host_outputs = host_outputs
|
||
self.cuda_outputs = cuda_outputs
|
||
self.bindings = bindings
|
||
self.batch_size = engine.max_batch_size
|
||
|
||
def __call__(self,img_np_nchw):
|
||
'''
|
||
TensorRT推理
|
||
:param img_np_nchw: 输入图像
|
||
'''
|
||
self.ctx.push()
|
||
|
||
# Restore
|
||
stream = self.stream
|
||
context = self.context
|
||
engine = self.engine
|
||
host_inputs = self.host_inputs
|
||
cuda_inputs = self.cuda_inputs
|
||
host_outputs = self.host_outputs
|
||
cuda_outputs = self.cuda_outputs
|
||
bindings = self.bindings
|
||
|
||
np.copyto(host_inputs[0], img_np_nchw.ravel())
|
||
cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
|
||
context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
|
||
cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
|
||
stream.synchronize()
|
||
self.ctx.pop()
|
||
return host_outputs[0]
|
||
|
||
|
||
def destroy(self):
|
||
# Remove any context from the top of the context stack, deactivating it.
|
||
self.ctx.pop()
|