Parse the ONNX graph and create the corresponding TensorRT network definition. :param onnx_path: The path to the ONNX graph to load. :param batch_size: Static batch size to build the engine with. :param dynamic_batch_size: Dynamic batch size to build the engine with,
(self, onnx_path, batch_size, dynamic_batch_size=None)
| 136 | self.parser = None |
| 137 | |
| 138 | def create_network(self, onnx_path, batch_size, dynamic_batch_size=None): |
| 139 | """ |
| 140 | Parse the ONNX graph and create the corresponding TensorRT network definition. |
| 141 | :param onnx_path: The path to the ONNX graph to load. |
| 142 | :param batch_size: Static batch size to build the engine with. |
| 143 | :param dynamic_batch_size: Dynamic batch size to build the engine with, if given, |
| 144 | batch_size is ignored, pass as a comma-separated string or int list as MIN,OPT,MAX |
| 145 | """ |
| 146 | network_flags = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) |
| 147 | |
| 148 | self.network = self.builder.create_network(network_flags) |
| 149 | self.parser = trt.OnnxParser(self.network, self.trt_logger) |
| 150 | |
| 151 | onnx_path = os.path.realpath(onnx_path) |
| 152 | with open(onnx_path, "rb") as f: |
| 153 | if not self.parser.parse(f.read()): |
| 154 | log.error("Failed to load ONNX file: {}".format(onnx_path)) |
| 155 | for error in range(self.parser.num_errors): |
| 156 | log.error(self.parser.get_error(error)) |
| 157 | sys.exit(1) |
| 158 | |
| 159 | log.info("Network Description") |
| 160 | |
| 161 | inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)] |
| 162 | profile = self.builder.create_optimization_profile() |
| 163 | dynamic_inputs = False |
| 164 | for input in inputs: |
| 165 | log.info("Input '{}' with shape {} and dtype {}".format(input.name, input.shape, input.dtype)) |
| 166 | if input.shape[0] == -1: |
| 167 | dynamic_inputs = True |
| 168 | if dynamic_batch_size: |
| 169 | if type(dynamic_batch_size) is str: |
| 170 | dynamic_batch_size = [int(v) for v in dynamic_batch_size.split(",")] |
| 171 | assert len(dynamic_batch_size) == 3 |
| 172 | min_shape = [dynamic_batch_size[0]] + list(input.shape[1:]) |
| 173 | opt_shape = [dynamic_batch_size[1]] + list(input.shape[1:]) |
| 174 | max_shape = [dynamic_batch_size[2]] + list(input.shape[1:]) |
| 175 | profile.set_shape(input.name, min_shape, opt_shape, max_shape) |
| 176 | log.info("Input '{}' Optimization Profile with shape MIN {} / OPT {} / MAX {}".format( |
| 177 | input.name, min_shape, opt_shape, max_shape)) |
| 178 | else: |
| 179 | shape = [batch_size] + list(input.shape[1:]) |
| 180 | profile.set_shape(input.name, shape, shape, shape) |
| 181 | log.info("Input '{}' Optimization Profile with shape {}".format(input.name, shape)) |
| 182 | if dynamic_inputs: |
| 183 | self.config.add_optimization_profile(profile) |
| 184 | |
| 185 | outputs = [self.network.get_output(i) for i in range(self.network.num_outputs)] |
| 186 | for output in outputs: |
| 187 | log.info("Output '{}' with shape {} and dtype {}".format(output.name, output.shape, output.dtype)) |
| 188 | |
| 189 | def set_mixed_precision(self): |
| 190 | """ |