309 lines
13 KiB
Python
309 lines
13 KiB
Python
import tensorrt as trt
|
|
import pycuda.driver as cuda
|
|
import pycuda.autoinit
|
|
import numpy as np
|
|
import torch
|
|
import traceback
|
|
import os
|
|
from PIL import Image
|
|
|
|
TRT_LOGGER = trt.Logger()
|
|
SKIP_ENGINE_MODEL_CHECK = True
|
|
|
|
def get_engine(engine_file_path):
|
|
if os.path.exists(engine_file_path):
|
|
print(f"Loading engine from file {engine_file_path}...")
|
|
with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
|
|
return runtime.deserialize_cuda_engine(f.read())
|
|
else:
|
|
print(f"No file named {engine_file_path}! Please check the input.")
|
|
return None
|
|
|
|
|
|
def numpy_to_torch_dtype(np_dtype):
|
|
mapping = {
|
|
np.float32: torch.float,
|
|
np.float64: torch.double,
|
|
np.float16: torch.half,
|
|
np.int32: torch.int32,
|
|
np.int64: torch.int64,
|
|
np.int16: torch.int16,
|
|
np.int8: torch.int8,
|
|
np.uint8: torch.uint8,
|
|
np.bool_: torch.bool
|
|
}
|
|
return mapping.get(np_dtype, None)
|
|
|
|
def match_shape(a, b):
|
|
if(len(a) == len(b)):
|
|
return tuple(a) == tuple(b)
|
|
elif len(a) > len(b):
|
|
if(a[0] == 1):
|
|
return match_shape(a[1:], b)
|
|
else:
|
|
if(b[0] == 1):
|
|
return match_shape(a, b[1:])
|
|
return False
|
|
|
|
def match_dtype(a, b):
|
|
if(a.__class__ == torch.dtype):
|
|
a = torch.tensor(0,dtype=a).numpy().dtype
|
|
return a == b
|
|
|
|
|
|
class EngineModel:
|
|
def __init__(self, engine_file_path, stream = None, device_int = 0, extra_lock = None):
|
|
self.device_int = device_int
|
|
self.extra_lock = extra_lock
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.acquire()
|
|
assert os.path.exists(engine_file_path), "Engine model path not exists!"
|
|
self.ctx = cuda.Device(self.device_int).make_context()
|
|
try:
|
|
self.engine = get_engine(engine_file_path) # 载入TensorRT引擎
|
|
input_nvars = 0
|
|
output_nvars = 0
|
|
self.input_names = []
|
|
self.output_names = []
|
|
|
|
# 【辅助函数】用于获取安全的 Shape (消除 -1)
|
|
def get_safe_shape(engine, name):
|
|
shape = engine.get_tensor_shape(name)
|
|
# 如果形状里包含 -1 (动态维度)
|
|
if -1 in shape:
|
|
# 获取 Profile 0 的 (min, opt, max)
|
|
# 取下标 [2] 即 Max Shape，确保分配足够大的显存
|
|
profile = engine.get_tensor_profile_shape(name, 0)
|
|
if profile:
|
|
print(f"[EngineModel] Detected dynamic shape for {name}: {shape} -> Using Max Profile: {profile[2]}")
|
|
return profile[2]
|
|
else:
|
|
# 如果获取不到 Profile (通常发生在 Output)，这是一个风险点
|
|
# 这里为了防止报错，可以尝试打印警告
|
|
print(f"[EngineModel] Warning: Dynamic output {name} has no profile. Mem alloc might fail.")
|
|
return shape
|
|
|
|
for binding in self.engine: # 遍历所有tensor，区分Input/Output
|
|
mode = self.engine.get_tensor_mode(binding)
|
|
if(mode== trt.TensorIOMode.INPUT):
|
|
input_nvars += 1
|
|
self.input_names.append(binding)
|
|
elif(mode == trt.TensorIOMode.OUTPUT):
|
|
output_nvars += 1
|
|
self.output_names.append(binding)
|
|
|
|
self.input_nvars = input_nvars # input的数量
|
|
self.output_nvars = output_nvars # output的数量
|
|
|
|
self.input_shapes = {name : get_safe_shape(self.engine, name) for name in self.input_names} # 获取每个 I/O 的 shape 和 dtype
|
|
self.input_dtypes = {name : self.engine.get_tensor_dtype(name) for name in self.input_names}
|
|
self.input_nbytes = {
|
|
name : trt.volume(self.input_shapes[name]) * trt.nptype(self.input_dtypes[name])().itemsize
|
|
for name in self.input_names
|
|
} # nbytes = tensor 占多少 CUDA 内存（字节数）
|
|
self.output_shapes = {name : get_safe_shape(self.engine, name) for name in self.output_names}
|
|
self.output_dtypes = {name : self.engine.get_tensor_dtype(name) for name in self.output_names}
|
|
self.output_nbytes = {
|
|
name : trt.volume(self.output_shapes[name]) * trt.nptype(self.output_dtypes[name])().itemsize
|
|
for name in self.output_names
|
|
}
|
|
self.dinputs = {name : cuda.mem_alloc(self.input_nbytes[name]) for name in self.input_names} # 为每个输入/输出分配 CUDA 设备内存
|
|
self.doutputs = {name :cuda.mem_alloc(self.output_nbytes[name]) for name in self.output_names}
|
|
self.context = self.engine.create_execution_context() # 创建 ExecutionContext（执行上下文）
|
|
if stream is None:
|
|
self.stream = cuda.Stream()
|
|
else:
|
|
self.stream = stream
|
|
for name in self.input_names: # 绑定 tensor 到 context
|
|
self.context.set_tensor_address(name, int(self.dinputs[name]))
|
|
for name in self.output_names:
|
|
self.context.set_tensor_address(name, int(self.doutputs[name]))
|
|
self.houtputs = {
|
|
name :
|
|
cuda.pagelocked_empty(
|
|
trt.volume(self.output_shapes[name]), dtype=trt.nptype(self.output_dtypes[name])
|
|
) for name in self.output_names
|
|
} # 分配 page-locked host 内存以存储输出
|
|
except:
|
|
self.ctx.pop()
|
|
raise Exception("CUDA Initialization Failed!")
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
|
|
def __call__(self, skip_check=SKIP_ENGINE_MODEL_CHECK, output_list=[], return_tensor=False, **inputs):
|
|
if not skip_check:
|
|
for name in inputs:
|
|
assert name in self.input_names
|
|
assert match_shape(inputs[name].shape, self.input_shapes[name])
|
|
assert match_dtype(inputs[name].dtype, trt.nptype(self.input_dtypes[name]))
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.acquire()
|
|
self.ctx.push()
|
|
r = {}
|
|
try:
|
|
|
|
for name in inputs:
|
|
hinput = inputs[name]
|
|
if (isinstance(hinput,torch.Tensor) and hinput.device.type=="cuda" and hinput.device.index==self.device_int):
|
|
hinput_con = hinput.contiguous()
|
|
ptr = hinput_con.data_ptr()
|
|
cuda.memcpy_dtod_async(self.dinputs[name], ptr, self.input_nbytes[name], self.stream)
|
|
else:
|
|
hinput_con = np.ascontiguousarray(hinput)
|
|
cuda.memcpy_htod_async(self.dinputs[name], hinput_con, self.stream)
|
|
for name in self.input_names:
|
|
if name not in inputs:
|
|
self.context.set_input_shape(name, self.input_shapes[name])
|
|
self.context.execute_async_v3(self.stream.handle)
|
|
if(return_tensor):
|
|
for name in output_list:
|
|
t = torch.zeros(trt.volume(self.output_shapes[name]), device=f"cuda:{self.device_int}", dtype=numpy_to_torch_dtype(trt.nptype(self.output_dtypes[name])))
|
|
ptr = t.data_ptr()
|
|
cuda.memcpy_dtod_async(ptr, self.doutputs[name], self.output_nbytes[name], self.stream)
|
|
t = t.reshape(tuple(self.output_shapes[name]))
|
|
r[name] = t
|
|
else:
|
|
for name in output_list:
|
|
cuda.memcpy_dtoh_async(self.houtputs[name], self.doutputs[name], self.stream)
|
|
r[name] = self.houtputs[name]
|
|
self.stream.synchronize()
|
|
except Exception as e:
|
|
print("TensorRT Execution Failed!")
|
|
traceback.print_exc()
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return None
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return r
|
|
|
|
|
|
def prefill(self, skip_check=SKIP_ENGINE_MODEL_CHECK, **inputs):
|
|
if not (skip_check):
|
|
for name in inputs:
|
|
in_input = (name in self.input_names)
|
|
assert in_input or (name in self.output_names)
|
|
assert match_shape(inputs[name].shape, self.input_shapes[name] if in_input else self.output_shapes[name])
|
|
assert match_dtype(inputs[name].dtype, trt.nptype(self.input_dtypes[name] if in_input else self.output_dtypes[name]))
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.acquire()
|
|
self.ctx.push()
|
|
try:
|
|
for name in inputs:
|
|
in_input = (name in self.input_names)
|
|
hinput = inputs[name]
|
|
|
|
dst_ptr = self.dinputs[name] if in_input else self.doutputs[name]
|
|
real_nbytes = 0
|
|
if isinstance(hinput, torch.Tensor):
|
|
real_nbytes = hinput.numel() * hinput.element_size()
|
|
else:
|
|
# 假设是 numpy
|
|
real_nbytes = hinput.nbytes
|
|
|
|
if (isinstance(hinput,torch.Tensor) and hinput.device.type=="cuda" and hinput.device.index==self.device_int):
|
|
hinput_con = hinput.contiguous()
|
|
ptr = hinput_con.data_ptr()
|
|
cuda.memcpy_dtod_async(dst_ptr, ptr, real_nbytes, self.stream)
|
|
else:
|
|
hinput_con = np.ascontiguousarray(hinput)
|
|
cuda.memcpy_htod_async(dst_ptr, hinput, self.stream)
|
|
self.stream.synchronize()
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return False
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return True
|
|
|
|
def __repr__(self):
|
|
r = "TensorRTEngineModel(\n\tInput=[\n"
|
|
for name in self.input_names:
|
|
r += f"\t\t{name}: \t{trt.nptype(self.input_dtypes[name]).__name__}{self.input_shapes[name]},\n"
|
|
r += "\t],Output=[\n"
|
|
for name in self.output_names:
|
|
r += f"\t\t{name}: \t{trt.nptype(self.output_dtypes[name]).__name__}{self.output_shapes[name]},\n"
|
|
r+="\t]\n)"
|
|
return r
|
|
|
|
def link(self, other, var_map, skip_check=SKIP_ENGINE_MODEL_CHECK):
|
|
assert self.device_int == other.device_int
|
|
if not (skip_check):
|
|
for source in var_map:
|
|
assert source in other.output_names
|
|
target = var_map[source]
|
|
assert target in self.input_names
|
|
assert match_shape(other.output_shapes[source], self.input_shapes[target])
|
|
assert match_dtype(other.output_dtypes[source], self.input_dtypes[target])
|
|
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.acquire()
|
|
self.ctx.push()
|
|
try:
|
|
for source in var_map:
|
|
target = var_map[source]
|
|
self.context.set_tensor_address(target, int(other.doutputs[source]))
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return False
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return True
|
|
|
|
def bind(self, var_map, skip_check=SKIP_ENGINE_MODEL_CHECK):
|
|
if not (skip_check):
|
|
for source in var_map:
|
|
assert source in self.output_names
|
|
target = var_map[source]
|
|
assert target in self.input_names
|
|
assert match_shape(self.output_shapes[source], self.input_shapes[target])
|
|
assert match_dtype(self.output_dtypes[source], self.input_dtypes[target])
|
|
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.acquire()
|
|
self.ctx.push()
|
|
try:
|
|
for source in var_map:
|
|
target = var_map[source]
|
|
self.context.set_tensor_address(target, int(self.doutputs[source]))
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return False
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return True
|
|
|
|
def unlink(self):
|
|
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.acquire()
|
|
self.ctx.push()
|
|
try:
|
|
for name in self.input_names:
|
|
self.context.set_tensor_address(name, int(self.dinputs[name]))
|
|
except:
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return False
|
|
self.ctx.pop()
|
|
if not(self.extra_lock is None):
|
|
self.extra_lock.release()
|
|
return True
|
|
|