An auxiliary class to implement running of TRT optimized engines
| 113 | |
| 114 | |
| 115 | class TRTEngine: |
| 116 | """ |
| 117 | An auxiliary class to implement running of TRT optimized engines |
| 118 | |
| 119 | """ |
| 120 | |
| 121 | def __init__(self, plan_path, logger=None): |
| 122 | """ |
| 123 | Loads serialized engine, creates execution context and activates it |
| 124 | Args: |
| 125 | plan_path: path to serialized TRT engine. |
| 126 | logger: optional logger object |
| 127 | """ |
| 128 | self.plan_path = plan_path |
| 129 | self.logger = logger or get_logger("monai.networks.trt_compiler") |
| 130 | self.logger.info(f"Loading TensorRT engine: {self.plan_path}") |
| 131 | self.engine = engine_from_bytes(bytes_from_path(self.plan_path)) |
| 132 | self.tensors = OrderedDict() |
| 133 | self.cuda_graph_instance = None # cuda graph |
| 134 | self.context = self.engine.create_execution_context() |
| 135 | self.input_names = [] |
| 136 | self.output_names = [] |
| 137 | self.dtypes = [] |
| 138 | self.cur_profile = 0 |
| 139 | self.input_table = {} |
| 140 | dtype_dict = trt_to_torch_dtype_dict() |
| 141 | for idx in range(self.engine.num_io_tensors): |
| 142 | binding = self.engine[idx] |
| 143 | if self.engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT: |
| 144 | self.input_names.append(binding) |
| 145 | elif self.engine.get_tensor_mode(binding) == trt.TensorIOMode.OUTPUT: |
| 146 | self.output_names.append(binding) |
| 147 | dtype = dtype_dict[self.engine.get_tensor_dtype(binding)] |
| 148 | self.dtypes.append(dtype) |
| 149 | self.logger.info( |
| 150 | f"Loaded TensorRT engine: {self.plan_path}.\nInputs: {self.input_names}\nOutputs: {self.output_names}" |
| 151 | ) |
| 152 | |
| 153 | def allocate_buffers(self, device): |
| 154 | """ |
| 155 | Allocates outputs to run TRT engine |
| 156 | Args: |
| 157 | device: GPU device to allocate memory on |
| 158 | """ |
| 159 | ctx = self.context |
| 160 | |
| 161 | for i, binding in enumerate(self.output_names): |
| 162 | shape = list(ctx.get_tensor_shape(binding)) |
| 163 | if binding not in self.tensors or list(self.tensors[binding].shape) != shape: |
| 164 | t = torch.empty(shape, dtype=self.dtypes[i], device=device).contiguous() |
| 165 | self.tensors[binding] = t |
| 166 | ctx.set_tensor_address(binding, t.data_ptr()) |
| 167 | |
| 168 | def set_inputs(self, feed_dict, stream): |
| 169 | """ |
| 170 | Sets input bindings for TRT engine according to feed_dict |
| 171 | Args: |
| 172 | feed_dict: a dictionary [str->Tensor] |
no outgoing calls
no test coverage detected
searching dependent graphs…