| import argparse | |
| from pathlib import Path | |
| from typing import List, Optional, Tuple, Union | |
| try: | |
| import tensorrt as trt | |
| except Exception: | |
| trt = None | |
| import warnings | |
| import numpy as np | |
| import torch | |
| warnings.filterwarnings(action='ignore', category=DeprecationWarning) | |
| class EngineBuilder: | |
| def __init__( | |
| self, | |
| checkpoint: Union[str, Path], | |
| opt_shape: Union[Tuple, List] = (1, 3, 640, 640), | |
| device: Optional[Union[str, int, torch.device]] = None) -> None: | |
| checkpoint = Path(checkpoint) if isinstance(checkpoint, | |
| str) else checkpoint | |
| assert checkpoint.exists() and checkpoint.suffix == '.onnx' | |
| if isinstance(device, str): | |
| device = torch.device(device) | |
| elif isinstance(device, int): | |
| device = torch.device(f'cuda:{device}') | |
| self.checkpoint = checkpoint | |
| self.opt_shape = np.array(opt_shape, dtype=np.float32) | |
| self.device = device | |
| def __build_engine(self, | |
| scale: Optional[List[List]] = None, | |
| fp16: bool = True, | |
| with_profiling: bool = True) -> None: | |
| logger = trt.Logger(trt.Logger.WARNING) | |
| trt.init_libnvinfer_plugins(logger, namespace='') | |
| builder = trt.Builder(logger) | |
| config = builder.create_builder_config() | |
| config.max_workspace_size = torch.cuda.get_device_properties( | |
| self.device).total_memory | |
| flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) | |
| network = builder.create_network(flag) | |
| parser = trt.OnnxParser(network, logger) | |
| if not parser.parse_from_file(str(self.checkpoint)): | |
| raise RuntimeError( | |
| f'failed to load ONNX file: {str(self.checkpoint)}') | |
| inputs = [network.get_input(i) for i in range(network.num_inputs)] | |
| outputs = [network.get_output(i) for i in range(network.num_outputs)] | |
| profile = None | |
| dshape = -1 in network.get_input(0).shape | |
| if dshape: | |
| profile = builder.create_optimization_profile() | |
| if scale is None: | |
| scale = np.array( | |
| [[1, 1, 0.5, 0.5], [1, 1, 1, 1], [4, 1, 1.5, 1.5]], | |
| dtype=np.float32) | |
| scale = (self.opt_shape * scale).astype(np.int32) | |
| elif isinstance(scale, List): | |
| scale = np.array(scale, dtype=np.int32) | |
| assert scale.shape[0] == 3, 'Input a wrong scale list' | |
| else: | |
| raise NotImplementedError | |
| for inp in inputs: | |
| logger.log( | |
| trt.Logger.WARNING, | |
| f'input "{inp.name}" with shape{inp.shape} {inp.dtype}') | |
| if dshape: | |
| profile.set_shape(inp.name, *scale) | |
| for out in outputs: | |
| logger.log( | |
| trt.Logger.WARNING, | |
| f'output "{out.name}" with shape{out.shape} {out.dtype}') | |
| if fp16 and builder.platform_has_fast_fp16: | |
| config.set_flag(trt.BuilderFlag.FP16) | |
| self.weight = self.checkpoint.with_suffix('.engine') | |
| if dshape: | |
| config.add_optimization_profile(profile) | |
| if with_profiling: | |
| config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED | |
| with builder.build_engine(network, config) as engine: | |
| self.weight.write_bytes(engine.serialize()) | |
| logger.log( | |
| trt.Logger.WARNING, f'Build tensorrt engine finish.\n' | |
| f'Save in {str(self.weight.absolute())}') | |
| def build(self, | |
| scale: Optional[List[List]] = None, | |
| fp16: bool = True, | |
| with_profiling=True): | |
| self.__build_engine(scale, fp16, with_profiling) | |
| def parse_args(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('checkpoint', help='Checkpoint file') | |
| parser.add_argument( | |
| '--img-size', | |
| nargs='+', | |
| type=int, | |
| default=[640, 640], | |
| help='Image size of height and width') | |
| parser.add_argument( | |
| '--device', type=str, default='cuda:0', help='TensorRT builder device') | |
| parser.add_argument( | |
| '--scales', | |
| type=str, | |
| default='[[1,3,640,640],[1,3,640,640],[1,3,640,640]]', | |
| help='Input scales for build dynamic input shape engine') | |
| parser.add_argument( | |
| '--fp16', action='store_true', help='Build model with fp16 mode') | |
| args = parser.parse_args() | |
| args.img_size *= 2 if len(args.img_size) == 1 else 1 | |
| return args | |
| def main(args): | |
| img_size = (1, 3, *args.img_size) | |
| try: | |
| scales = eval(args.scales) | |
| except Exception: | |
| print('Input scales is not a python variable') | |
| print('Set scales default None') | |
| scales = None | |
| builder = EngineBuilder(args.checkpoint, img_size, args.device) | |
| builder.build(scales, fp16=args.fp16) | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| main(args) | |