| 52 | |
| 53 | |
| 54 | class EngineModel: |
| 55 | def __init__(self, engine_file_path, stream = None, device_int = 0, extra_lock = None): |
| 56 | self.device_int = device_int |
| 57 | self.extra_lock = extra_lock |
| 58 | if not(self.extra_lock is None): |
| 59 | self.extra_lock.acquire() |
| 60 | assert os.path.exists(engine_file_path), "Engine model path not exists!" |
| 61 | self.ctx = cuda.Device(self.device_int).make_context() |
| 62 | try: |
| 63 | self.engine = get_engine(engine_file_path) |
| 64 | input_nvars = 0 |
| 65 | output_nvars = 0 |
| 66 | self.input_names = [] |
| 67 | self.output_names = [] |
| 68 | for binding in self.engine: |
| 69 | mode = self.engine.get_tensor_mode(binding) |
| 70 | if(mode== trt.TensorIOMode.INPUT): |
| 71 | input_nvars += 1 |
| 72 | self.input_names.append(binding) |
| 73 | elif(mode == trt.TensorIOMode.OUTPUT): |
| 74 | output_nvars += 1 |
| 75 | self.output_names.append(binding) |
| 76 | |
| 77 | self.input_nvars = input_nvars |
| 78 | self.output_nvars = output_nvars |
| 79 | |
| 80 | self.input_shapes = {name : self.engine.get_tensor_shape(name) for name in self.input_names} |
| 81 | self.input_dtypes = {name : self.engine.get_tensor_dtype(name) for name in self.input_names} |
| 82 | self.input_nbytes = { |
| 83 | name : trt.volume(self.input_shapes[name]) * trt.nptype(self.input_dtypes[name])().itemsize |
| 84 | for name in self.input_names |
| 85 | } |
| 86 | self.output_shapes = {name : self.engine.get_tensor_shape(name) for name in self.output_names} |
| 87 | self.output_dtypes = {name : self.engine.get_tensor_dtype(name) for name in self.output_names} |
| 88 | self.output_nbytes = { |
| 89 | name : trt.volume(self.output_shapes[name]) * trt.nptype(self.output_dtypes[name])().itemsize |
| 90 | for name in self.output_names |
| 91 | } |
| 92 | self.dinputs = {name : cuda.mem_alloc(self.input_nbytes[name]) for name in self.input_names} |
| 93 | self.doutputs = {name :cuda.mem_alloc(self.output_nbytes[name]) for name in self.output_names} |
| 94 | self.context = self.engine.create_execution_context() |
| 95 | if stream is None: |
| 96 | self.stream = cuda.Stream() |
| 97 | else: |
| 98 | self.stream = stream |
| 99 | for name in self.input_names: |
| 100 | self.context.set_tensor_address(name, int(self.dinputs[name])) |
| 101 | for name in self.output_names: |
| 102 | self.context.set_tensor_address(name, int(self.doutputs[name])) |
| 103 | self.houtputs = { |
| 104 | name : |
| 105 | cuda.pagelocked_empty( |
| 106 | trt.volume(self.output_shapes[name]), dtype=trt.nptype(self.output_dtypes[name]) |
| 107 | ) for name in self.output_names |
| 108 | } |
| 109 | except: |
| 110 | self.ctx.pop() |
| 111 | raise Exception("CUDA Initialization Failed!") |