使用 trt 的int8 量化和推断 onnx 模型

2022-04-06 更新：

理清几个概念：

1、onnx模型本身要有动态维度，否则只能转静态维度的trt engine。

2、只要一个profile就够了，设个最小最大维度，最优就是最常用的维度。在推断的时候要绑定一下。

3、builder 和 config 里有很多相同的设置，如果用了 config，就不需要设置 builder中的相同参数了。

def onnx_2_trt(onnx_filename, engine_filename, mode='fp32', max_batch_size=1, min_wh=(160,160), max_wh=(320,320), int8_calib=None):
    ''' convert onnx to tensorrt engine, use mode of ['fp32', 'fp16', 'int8']
    :return: trt engine
    '''

    assert mode in ['fp32', 'fp16', 'int8'], "mode should be in ['fp32', 'fp16', 'int8']"

    G_LOGGER = trt.Logger(trt.Logger.WARNING)
    # TRT7中的onnx解析器的network，需要指定EXPLICIT_BATCH
    EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    with trt.Builder(G_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, \
            trt.OnnxParser(network, G_LOGGER) as parser:

        print('Loading ONNX file from path %s...'%(onnx_filename))
        with open(onnx_filename, 'rb') as model:
            print('Beginning ONNX file parsing')
            if not parser.parse(model.read()):
                for e in range(parser.num_errors):
                    print(parser.get_error(e))
                raise TypeError('Parser parse failed.')

        print('Completed parsing of ONNX file')
        
        # wujp 2022-03-29 如果使用了config，builder就不用设置了。
        builder.max_batch_size = max_batch_size # max_batch_size 在config中没有
        #builder.max_workspace_size = 1 << 30
            
        if mode == 'int8':
            assert (builder.platform_has_fast_int8 == True), 'not support int8'
            #builder.int8_mode = True
            #builder.int8_calibrator = calib
        elif mode == 'fp16':
            assert (builder.platform_has_fast_fp16 == True), 'not support fp16'
            #builder.fp16_mode = True

        profile = builder.create_optimization_profile()
        inputs = [network.get_input(i) for i in range(network.num_inputs)]
        for inp in inputs:
            fbs, shape = inp.shape[0], inp.shape[1:]
            if (shape[1] == -1 and shape[2] == -1): # height 和 width 都是动态的。
                profile.set_shape(inp.name, min=(1, shape[0], *min_wh), opt=(8, shape[0], *min_wh), max=(max_batch_size, shape[0], *max_wh))
            else:
                profile.set_shape(inp.name, min=(1, *shape), opt=(8, *shape), max=(max_batch_size, *shape))
        
        config = builder.create_builder_config()
        config.max_workspace_size = 1 << 30
        if mode == 'int8':
            config.set_flag(trt.BuilderFlag.INT8)
            config.int8_calibrator = int8_calib
        elif mode == 'fp16':
            config.set_flag(trt.BuilderFlag.FP16)
        config.add_optimization_profile(profile)
        config.set_calibration_profile(profile) # 不加会有警告 [TensorRT] WARNING: Calibration Profile is not defined. Runing calibration with Profile 0
        print('Building an engine from file %s; this may take a while...'%(onnx_filename))
        #engine = builder.build_cuda_engine(network, config)
        engine = builder.build_engine(network, config)
        print('Created engine success! ')

        # 保存计划文件
        print('Saving TRT engine file to path %s...'%(engine_filename))
        with open(engine_filename, 'wb') as f:
            f.write(engine.serialize())
        print('Engine file has already saved to %s!'%(engine_filename))
        return engine

--------------------------------------------------------------------

2022-03-27

以下代码单batch没问题，多batch不行。

目录

生成 trt 模型