| 634 | |
| 635 | class TRT: |
| 636 | def __init__(self, engine_path: Path): |
| 637 | import tensorrt as trt |
| 638 | |
| 639 | trt.init_libnvinfer_plugins(None, "") |
| 640 | |
| 641 | self._runtime = trt.Runtime(trt.Logger(trt.Logger.ERROR)) |
| 642 | self._engine = self._runtime.deserialize_cuda_engine(engine_path.read_bytes()) |
| 643 | |
| 644 | if self._engine is None: |
| 645 | raise RuntimeError( |
| 646 | f"Failed to deserialize TensorRT engine from {engine_path}" |
| 647 | ) |
| 648 | |
| 649 | self._context = self._engine.create_execution_context() |
| 650 | |
| 651 | # Allocate the output tensors (inputs are dynamically setup) |
| 652 | self._input_names = [] |
| 653 | self._output_names = [] |
| 654 | self._output_tensors = [] |
| 655 | for i in range(self._engine.num_io_tensors): |
| 656 | tensor_name = self._engine.get_tensor_name(i) |
| 657 | if self._engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.OUTPUT: |
| 658 | self._output_names.append(tensor_name) |
| 659 | shape = self._engine.get_tensor_shape(tensor_name) |
| 660 | |
| 661 | # Determine dtype based on tensor name |
| 662 | if "labels" in tensor_name or "num_detections" in tensor_name: |
| 663 | dtype = cvcuda.Type.S32 |
| 664 | else: |
| 665 | dtype = cvcuda.Type.F32 |
| 666 | |
| 667 | # contiguous_tensor = cvcuda.Tensor((np.prod(shape),), dtype=dtype) |
| 668 | # self._output_tensors.append(contiguous_tensor.reshape(tuple(shape))) |
| 669 | self._output_tensors.append(cvcuda.Tensor(tuple(shape), dtype)) |
| 670 | else: |
| 671 | self._input_names.append(tensor_name) |
| 672 | |
| 673 | # Set the output tensor addresses |
| 674 | for i, tensor in enumerate(self._output_tensors): |
| 675 | self._context.set_tensor_address( |
| 676 | self._output_names[i], tensor.cuda().__cuda_array_interface__["data"][0] |
| 677 | ) |
| 678 | |
| 679 | def __call__(self, tensors: list[cvcuda.Tensor]) -> list[cvcuda.Tensor]: |
| 680 | for i, tensor in enumerate(tensors): |