Parses an ONNX graph and builds a TensorRT engine from it.
| 113 | |
| 114 | |
| 115 | class EngineBuilder: |
| 116 | """ |
| 117 | Parses an ONNX graph and builds a TensorRT engine from it. |
| 118 | """ |
| 119 | |
| 120 | def __init__(self, verbose=False, workspace=8): |
| 121 | """ |
| 122 | :param verbose: If enabled, a higher verbosity level will be set on the TensorRT logger. |
| 123 | :param workspace: Max memory workspace to allow, in Gb. |
| 124 | """ |
| 125 | self.trt_logger = trt.Logger(trt.Logger.INFO) |
| 126 | if verbose: |
| 127 | self.trt_logger.min_severity = trt.Logger.Severity.VERBOSE |
| 128 | |
| 129 | trt.init_libnvinfer_plugins(self.trt_logger, namespace="") |
| 130 | |
| 131 | self.builder = trt.Builder(self.trt_logger) |
| 132 | self.config = self.builder.create_builder_config() |
| 133 | self.config.max_workspace_size = workspace * (2 ** 30) |
| 134 | |
| 135 | self.batch_size = None |
| 136 | self.network = None |
| 137 | self.parser = None |
| 138 | |
| 139 | def create_network(self, onnx_path): |
| 140 | """ |
| 141 | Parse the ONNX graph and create the corresponding TensorRT network definition. |
| 142 | :param onnx_path: The path to the ONNX graph to load. |
| 143 | """ |
| 144 | network_flags = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) |
| 145 | |
| 146 | self.network = self.builder.create_network(network_flags) |
| 147 | self.parser = trt.OnnxParser(self.network, self.trt_logger) |
| 148 | |
| 149 | onnx_path = os.path.realpath(onnx_path) |
| 150 | with open(onnx_path, "rb") as f: |
| 151 | if not self.parser.parse(f.read()): |
| 152 | log.error("Failed to load ONNX file: {}".format(onnx_path)) |
| 153 | for error in range(self.parser.num_errors): |
| 154 | log.error(self.parser.get_error(error)) |
| 155 | sys.exit(1) |
| 156 | |
| 157 | inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)] |
| 158 | outputs = [self.network.get_output(i) for i in range(self.network.num_outputs)] |
| 159 | |
| 160 | log.info("Network Description") |
| 161 | for input in inputs: |
| 162 | self.batch_size = input.shape[0] |
| 163 | log.info("Input '{}' with shape {} and dtype {}".format(input.name, input.shape, input.dtype)) |
| 164 | for output in outputs: |
| 165 | log.info("Output '{}' with shape {} and dtype {}".format(output.name, output.shape, output.dtype)) |
| 166 | assert self.batch_size > 0 |
| 167 | self.builder.max_batch_size = self.batch_size |
| 168 | |
| 169 | def create_engine(self, engine_path, precision, config_file, calib_input=None, calib_cache=None, calib_num_images=5000, |
| 170 | calib_batch_size=8): |
| 171 | """ |
| 172 | Build the TensorRT engine and serialize it to disk. |