Spaces:
No application file
No application file
File size: 4,225 Bytes
f113e60 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
# import tensorrt as trt
# import pycuda.driver as cuda
# class HostDeviceMem(object):
# def __init__(self, host_mem, device_mem) -> None:
# self.host = host_mem
# self.device = device_mem
# def __str__(self) -> str:
# return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
# def __repr__(self):
# return self.__str__()
# class TensorrtBase:
# def __init__(self, engine_file_path, input_names, output_names, *, gpu_id=0, dynamic_factor=1, max_batch_size=1) -> None:
# self.input_names = input_names
# self.output_names = output_names
# self.trt_logger = trt.Logger(trt.Logger.WARNING)
# self.cuda_ctx = cuda.Device(gpu_id).make_context()
# self.max_batch_size = max_batch_size
# self.engine = self._load_engine(engine_file_path)
# self.binding_names = self.input_names + self.output_names
# self.context = self.engine.create_execution_context()
# self.buffers = self._allocate_buffer(dynamic_factor)
# def _load_engine(self, engine_file_path):
# # Force init TensorRT plugins
# trt.init_libnvinfer_plugins(None, '')
# with open(engine_file_path, "rb") as f, \
# trt.Runtime(self.trt_logger) as runtime:
# engine = runtime.deserialize_cuda_engine(f.read())
# return engine
# def _allocate_buffer(self, dynamic_factor):
# """Allocate buffer
# :dynamic_factor: normally expand the buffer size for dynamic shape
# """
# inputs = []
# outputs = []
# bindings = [None] * len(self.binding_names)
# stream = cuda.Stream()
# for binding in self.binding_names:
# binding_idx = self.engine[binding]
# if binding_idx == -1:
# print("❌ Binding Names!")
# continue
# # trt.volume() return negtive volue if -1 in shape
# size = abs(trt.volume(self.engine.get_binding_shape(binding))) * \
# self.max_batch_size * dynamic_factor
# dtype = trt.nptype(self.engine.get_binding_dtype(binding))
# # Allocate host and device buffers
# host_mem = cuda.pagelocked_empty(size, dtype)
# device_mem = cuda.mem_alloc(host_mem.nbytes)
# # Append the device buffer to device bindings.
# bindings[binding_idx] = int(device_mem)
# # Append to the appropriate list.
# if self.engine.binding_is_input(binding):
# inputs.append(HostDeviceMem(host_mem, device_mem))
# else:
# outputs.append(HostDeviceMem(host_mem, device_mem))
# return inputs, outputs, bindings, stream
# # def do_inference(self, inf_in_list, *, binding_shape_map=None):
# # """Main function for inference
# # :inf_in_list: input list.
# # :binding_shape_map: {<binding_name>: <shape>}, leave it to None for fixed shape
# # """
# # inputs, outputs, bindings, stream = self.buffers
# # if binding_shape_map:
# # self.context.active_optimization_profile = 0
# # for binding_name, shape in binding_shape_map.items():
# # binding_idx = self.engine[binding_name]
# # self.context.set_binding_shape(binding_idx, shape)
# # # transfer input data to device
# # for i in range(len(inputs)):
# # inputs[i].host = inf_in_list[i]
# # cuda.memcpy_htod_async(inputs[i].device, inputs[i].host, stream)
# # # do inference
# # # context.profiler = trt.Profiler()
# # self.context.execute_async_v2(bindings=bindings,
# # stream_handle=stream.handle)
# # # copy data from device to host
# # for i in range(len(outputs)):
# # cuda.memcpy_dtoh_async(outputs[i].host, outputs[i].device, stream)
# # stream.synchronize()
# # trt_outputs = [out.host.copy() for out in outputs]
# # return trt_outputs
# def __del__(self):
# self.cuda_ctx.pop()
# del self.cuda_ctx |