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', or 'mixed'. :param calib_input: The path to a directory holding the
(self, engine_path, precision, calib_input=None, calib_cache=None, calib_num_images=5000,
calib_batch_size=8)
| 217 | log.info("Mixed-Precision Layer {} set to HALF STRICT data type".format(layer.name)) |
| 218 | |
| 219 | def create_engine(self, engine_path, precision, calib_input=None, calib_cache=None, calib_num_images=5000, |
| 220 | calib_batch_size=8): |
| 221 | """ |
| 222 | Build the TensorRT engine and serialize it to disk. |
| 223 | :param engine_path: The path where to serialize the engine to. |
| 224 | :param precision: The datatype to use for the engine, either 'fp32', 'fp16', 'int8', or 'mixed'. |
| 225 | :param calib_input: The path to a directory holding the calibration images. |
| 226 | :param calib_cache: The path where to write the calibration cache to, or if it already exists, load it from. |
| 227 | :param calib_num_images: The maximum number of images to use for calibration. |
| 228 | :param calib_batch_size: The batch size to use for the calibration process. |
| 229 | """ |
| 230 | engine_path = os.path.realpath(engine_path) |
| 231 | engine_dir = os.path.dirname(engine_path) |
| 232 | os.makedirs(engine_dir, exist_ok=True) |
| 233 | log.info("Building {} Engine in {}".format(precision, engine_path)) |
| 234 | |
| 235 | inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)] |
| 236 | |
| 237 | if precision in ["fp16", "int8", "mixed"]: |
| 238 | if not self.builder.platform_has_fast_fp16: |
| 239 | log.warning("FP16 is not supported natively on this platform/device") |
| 240 | self.config.set_flag(trt.BuilderFlag.FP16) |
| 241 | if precision in ["int8", "mixed"]: |
| 242 | if not self.builder.platform_has_fast_int8: |
| 243 | log.warning("INT8 is not supported natively on this platform/device") |
| 244 | self.config.set_flag(trt.BuilderFlag.INT8) |
| 245 | self.config.int8_calibrator = EngineCalibrator(calib_cache) |
| 246 | if calib_cache is None or not os.path.exists(calib_cache): |
| 247 | calib_shape = [calib_batch_size] + list(inputs[0].shape[1:]) |
| 248 | calib_dtype = trt.nptype(inputs[0].dtype) |
| 249 | self.config.int8_calibrator.set_image_batcher( |
| 250 | ImageBatcher(calib_input, calib_shape, calib_dtype, max_num_images=calib_num_images, |
| 251 | exact_batches=True, shuffle_files=True)) |
| 252 | |
| 253 | engine_bytes = None |
| 254 | try: |
| 255 | engine_bytes = self.builder.build_serialized_network(self.network, self.config) |
| 256 | except AttributeError: |
| 257 | engine = self.builder.build_engine(self.network, self.config) |
| 258 | engine_bytes = engine.serialize() |
| 259 | del engine |
| 260 | assert engine_bytes |
| 261 | with open(engine_path, "wb") as f: |
| 262 | log.info("Serializing engine to file: {:}".format(engine_path)) |
| 263 | f.write(engine_bytes) |
| 264 | |
| 265 | |
| 266 | def main(args): |
no test coverage detected