Spaces:
Running
Running
| 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 | |