(self, skip_check=SKIP_ENGINE_MODEL_CHECK, **inputs)
| 161 | |
| 162 | |
| 163 | def prefill(self, skip_check=SKIP_ENGINE_MODEL_CHECK, **inputs): |
| 164 | if not (skip_check): |
| 165 | for name in inputs: |
| 166 | in_input = (name in self.input_names) |
| 167 | assert in_input or (name in self.output_names) |
| 168 | assert match_shape(inputs[name].shape, self.input_shapes[name] if in_input else self.output_shapes[name]) |
| 169 | assert match_dtype(inputs[name].dtype, trt.nptype(self.input_dtypes[name] if in_input else self.output_dtypes[name])) |
| 170 | if not(self.extra_lock is None): |
| 171 | self.extra_lock.acquire() |
| 172 | self.ctx.push() |
| 173 | try: |
| 174 | for name in inputs: |
| 175 | in_input = (name in self.input_names) |
| 176 | hinput = inputs[name] |
| 177 | if (isinstance(hinput,torch.Tensor) and hinput.device.type=="cuda" and hinput.device.index==self.device_int): |
| 178 | hinput_con = hinput.contiguous() |
| 179 | ptr = hinput_con.data_ptr() |
| 180 | cuda.memcpy_dtod_async(self.dinputs[name] if in_input else self.doutputs[name], ptr, self.input_nbytes[name] if in_input else self.output_nbytes[name], self.stream) |
| 181 | else: |
| 182 | hinput_con = np.ascontiguousarray(hinput) |
| 183 | cuda.memcpy_htod_async(self.dinputs[name] if in_input else self.doutputs[name], hinput, self.stream) |
| 184 | self.stream.synchronize() |
| 185 | except Exception as e: |
| 186 | traceback.print_exc() |
| 187 | self.ctx.pop() |
| 188 | if not(self.extra_lock is None): |
| 189 | self.extra_lock.release() |
| 190 | return False |
| 191 | self.ctx.pop() |
| 192 | if not(self.extra_lock is None): |
| 193 | self.extra_lock.release() |
| 194 | return True |
| 195 | |
| 196 | def __repr__(self): |
| 197 | r = "TensorRTEngineModel(\n\tInput=[\n" |
no test coverage detected