@brief: Run the TensorRT engine with the given inputs and outputs @param inputs: dict of input tensors, key is tensor name, value is tensor pointer or torch tensor @param outputs: dict of output tensors, key is tensor name, value is tensor pointer or torch tensor @pa
(self,
inputs: Dict[str, Any],
outputs: Dict[str, Any],
stream,
context=None)
| 266 | raise |
| 267 | |
| 268 | def run(self, |
| 269 | inputs: Dict[str, Any], |
| 270 | outputs: Dict[str, Any], |
| 271 | stream, |
| 272 | context=None) -> bool: |
| 273 | ''' |
| 274 | @brief: Run the TensorRT engine with the given inputs and outputs |
| 275 | @param inputs: dict of input tensors, key is tensor name, value is tensor pointer or torch tensor |
| 276 | @param outputs: dict of output tensors, key is tensor name, value is tensor pointer or torch tensor |
| 277 | @param stream: cuda stream to enqueue the TensorRT engine on |
| 278 | @param context: TensorRT execution context, if None, use the default context |
| 279 | @return: True if enqueue succeeded, note the enqueue is an async call, |
| 280 | returning True does not mean the execution is finished |
| 281 | ''' |
| 282 | # enqueue to the default context if context is not specified |
| 283 | if context is None: |
| 284 | context = self.context |
| 285 | |
| 286 | import torch |
| 287 | for tensor_name in inputs: |
| 288 | tensor = inputs[tensor_name] |
| 289 | ptr = tensor.data_ptr() if isinstance(tensor, |
| 290 | torch.Tensor) else tensor |
| 291 | context.set_tensor_address(tensor_name, ptr) |
| 292 | for tensor_name in outputs: |
| 293 | tensor = outputs[tensor_name] |
| 294 | ptr = tensor.data_ptr() if isinstance(tensor, |
| 295 | torch.Tensor) else tensor |
| 296 | context.set_tensor_address(tensor_name, ptr) |
| 297 | ok = context.execute_async_v3(stream) |
| 298 | return ok |
| 299 | |
| 300 | def _debug_run(self, |
| 301 | inputs: Dict[str, "torch.Tensor"], |