Build the TensorRT engine and serialize it to disk. :param engine_path: The path where to serialize the engine to. :param precision: The datatype to use for the engine, either 'fp32', 'fp16', 'int8'. :param calib_input: The path to a directory holding the calibration
(self, engine_path, precision, config_file, calib_input=None, calib_cache=None, calib_num_images=5000,
calib_batch_size=8)
| 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. |
| 173 | :param engine_path: The path where to serialize the engine to. |
| 174 | :param precision: The datatype to use for the engine, either 'fp32', 'fp16', 'int8'. |
| 175 | :param calib_input: The path to a directory holding the calibration images. |
| 176 | :param calib_cache: The path where to write the calibration cache to, or if it already exists, load it from. |
| 177 | :param calib_num_images: The maximum number of images to use for calibration. |
| 178 | :param calib_batch_size: The batch size to use for the calibration process. |
| 179 | """ |
| 180 | engine_path = os.path.realpath(engine_path) |
| 181 | engine_dir = os.path.dirname(engine_path) |
| 182 | os.makedirs(engine_dir, exist_ok=True) |
| 183 | log.info("Building {} Engine in {}".format(precision, engine_path)) |
| 184 | |
| 185 | inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)] |
| 186 | |
| 187 | if precision in ["fp16", "int8"]: |
| 188 | if not self.builder.platform_has_fast_fp16: |
| 189 | log.warning("FP16 is not supported natively on this platform/device") |
| 190 | self.config.set_flag(trt.BuilderFlag.FP16) |
| 191 | if precision in ["int8"]: |
| 192 | if not self.builder.platform_has_fast_int8: |
| 193 | log.warning("INT8 is not supported natively on this platform/device") |
| 194 | self.config.set_flag(trt.BuilderFlag.INT8) |
| 195 | self.config.int8_calibrator = EngineCalibrator(calib_cache) |
| 196 | if calib_cache is None or not os.path.exists(calib_cache): |
| 197 | calib_shape = [calib_batch_size] + list(inputs[0].shape[1:]) |
| 198 | calib_dtype = trt.nptype(inputs[0].dtype) |
| 199 | self.config.int8_calibrator.set_image_batcher( |
| 200 | ImageBatcher(calib_input, calib_shape, calib_dtype, max_num_images=calib_num_images, |
| 201 | exact_batches=True, config_file=config_file)) |
| 202 | |
| 203 | engine_bytes = None |
| 204 | try: |
| 205 | engine_bytes = self.builder.build_serialized_network(self.network, self.config) |
| 206 | except AttributeError: |
| 207 | engine = self.builder.build_engine(self.network, self.config) |
| 208 | engine_bytes = engine.serialize() |
| 209 | del engine |
| 210 | assert engine_bytes |
| 211 | with open(engine_path, "wb") as f: |
| 212 | log.info("Serializing engine to file: {:}".format(engine_path)) |
| 213 | f.write(engine_bytes) |
| 214 | |
| 215 | |
| 216 | def main(args): |
no test coverage detected