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