Function build_engine

ppq/utils/TensorRTUtil.py:212–263

(
    onnx_file: str, engine_file: str,
    fp16: bool = True, int8: bool = False, 
    int8_scale_file: str = None,
    explicit_batch: bool = True, 
    workspace: int = 4294967296, # 4GB
    )

def build_engine(
    onnx_file: str, engine_file: str,
    fp16: bool = True, int8: bool = False, 
    int8_scale_file: str = None,
    explicit_batch: bool = True, 
    workspace: int = 4294967296, # 4GB
    ):
    TRT_LOGGER = trt.Logger()
    """
    Build a TensorRT Engine with given onnx model.

    Flag int8, fp16 specifies the precision of layer:
        For building FP32 engine: set int8 = False, fp16 = False, int8_scale_file = None
        For building FP16 engine: set int8 = False, fp16 = True, int8_scale_file = None
        For building INT8 engine: set True = False, fp16 = True, int8_scale_file = 'json file name'

    """

    if int8 is True:
        if int8_scale_file is None:
            raise ValueError('Build Quantized TensorRT Engine Requires a JSON file which specifies variable scales, '
                             'however int8_scale_file is None now.')

    builder = trt.Builder(TRT_LOGGER)
    if explicit_batch:
        network = builder.create_network(1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    else: network = builder.create_network()
    
    config = builder.create_builder_config()

    parser = trt.OnnxParser(network, TRT_LOGGER)
    config.max_workspace_size = workspace

    if not os.path.exists(onnx_file):
        raise FileNotFoundError(f'ONNX file {onnx_file} not found')

    with open(onnx_file, 'rb') as model:
        if not parser.parse(model.read()):
            print('ERROR: Failed to parse the ONNX file.')
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    if fp16: config.set_flag(trt.BuilderFlag.FP16)
    if int8_scale_file is not None and int8:
        config.set_flag(trt.BuilderFlag.INT8)
        setDynamicRange(network, int8_scale_file)

    engine = builder.build_engine(network, config)

    with open(engine_file, "wb") as f:
        f.write(engine.serialize())


class MyProfiler(trt.IProfiler):