(self, skip_check=SKIP_ENGINE_MODEL_CHECK, output_list=[], return_tensor=False, **inputs)
| 114 | self.extra_lock.release() |
| 115 | |
| 116 | def __call__(self, skip_check=SKIP_ENGINE_MODEL_CHECK, output_list=[], return_tensor=False, **inputs): |
| 117 | if not skip_check: |
| 118 | for name in inputs: |
| 119 | assert name in self.input_names |
| 120 | assert match_shape(inputs[name].shape, self.input_shapes[name]) |
| 121 | assert match_dtype(inputs[name].dtype, trt.nptype(self.input_dtypes[name])) |
| 122 | if not(self.extra_lock is None): |
| 123 | self.extra_lock.acquire() |
| 124 | self.ctx.push() |
| 125 | r = {} |
| 126 | try: |
| 127 | |
| 128 | for name in inputs: |
| 129 | hinput = inputs[name] |
| 130 | if (isinstance(hinput,torch.Tensor) and hinput.device.type=="cuda" and hinput.device.index==self.device_int): |
| 131 | hinput_con = hinput.contiguous() |
| 132 | ptr = hinput_con.data_ptr() |
| 133 | cuda.memcpy_dtod_async(self.dinputs[name], ptr, self.input_nbytes[name], self.stream) |
| 134 | else: |
| 135 | hinput_con = np.ascontiguousarray(hinput) |
| 136 | cuda.memcpy_htod_async(self.dinputs[name], hinput_con, self.stream) |
| 137 | self.context.execute_async_v3(self.stream.handle) |
| 138 | if(return_tensor): |
| 139 | for name in output_list: |
| 140 | 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]))) |
| 141 | ptr = t.data_ptr() |
| 142 | cuda.memcpy_dtod_async(ptr, self.doutputs[name], self.output_nbytes[name], self.stream) |
| 143 | t = t.reshape(tuple(self.output_shapes[name])) |
| 144 | r[name] = t |
| 145 | else: |
| 146 | for name in output_list: |
| 147 | cuda.memcpy_dtoh_async(self.houtputs[name], self.doutputs[name], self.stream) |
| 148 | r[name] = self.houtputs[name] |
| 149 | self.stream.synchronize() |
| 150 | except Exception as e: |
| 151 | print("TensorRT Execution Failed!") |
| 152 | traceback.print_exc() |
| 153 | self.ctx.pop() |
| 154 | if not(self.extra_lock is None): |
| 155 | self.extra_lock.release() |
| 156 | return None |
| 157 | self.ctx.pop() |
| 158 | if not(self.extra_lock is None): |
| 159 | self.extra_lock.release() |
| 160 | return r |
| 161 | |
| 162 | |
| 163 | def prefill(self, skip_check=SKIP_ENGINE_MODEL_CHECK, **inputs): |
nothing calls this directly
no test coverage detected