Tensorrt Bert Pytorch

基于Tensorrt 实现Bert的推理加速，记录一下流程和踩得坑

1.Tensorrt的安装

Tensorrt安装时要对应版本，我的版本：

tensorrt官网下载链接: nVidia.
安装时按照下面的两个教程按照即可：
1. 2.

2.pth转onnx文件

https://zhuanlan.zhihu.com/p/446477075 此文章中八1.内容写的非常详细。
需要注意的是，要看自己有几个输入，几个输出，动态轴如何设置就可以了

def export_dynamic_onnx_model(model, inputs, dynamic_onnx_model_path):
    if os.path.exists(dynamic_onnx_model_path):
        return
    symbolic_names = {0: 'batch_size', 1: 'max_seq_len'} 
    with torch.no_grad():
        torch.onnx.export(model, 
                        (inputs['input_ids'], inputs['attention_mask'], inputs['token_type_ids']),
                        dynamic_onnx_model_path, 
                        verbose=True, do_constant_folding=True,
                        input_names=['input_ids', 'input_mask', 'input_token_type_ids'], 
                        output_names=['output'], opset_version=13,
                        dynamic_axes={'input_ids':symbolic_names, 'input_mask':symbolic_names, 
                        'input_token_type_ids':symbolic_names}) # , 'output':symbolic_names
    logger.info("ONNX Model exported to {0}".format(dynamic_onnx_model_path))

其中opset_version默认是9，我的原网络代码中有Nonzero操作，此版本不支持，改成13就好

https://github.com/onnx/onnx/blob/main/docs/Operators.md 从这里可以看哪个版本支持哪些算子

https://blog.csdn.net/github_28260175/article/details/103436020 这篇文章的评论里也有设计到， 同时这篇文章提到使用onnx-tensorrt转trt时的一些坑，以及使用trt python api的解决办法，但比较难。

安装onnx-tensorrt的教程：https://blog.csdn.net/qq_36915686/article/details/120107756 里面也有好多坑，版本一定要对，看开头

好像使用torch2trt也可以实现推理 https://zhuanlan.zhihu.com/p/398451694

3.onnx转trt

①使用tensorrt自带的trtexec命令得到Trt文件

https://zhuanlan.zhihu.com/p/446477075有详细使用方法。
刚开始转的时候，报缺少nonzero算子的问题，
我的解决办法是像上面所说，在pth转onnx时，修改optset_version,找出支持nonzero算子的版本。

②动态设置batch_size

参考https://zhuanlan.zhihu.com/p/446477075的八的内容。
（pth转onnx时设置动态轴batch_size, 在./trtexec时要指定shape（batch_size=1时应该不需要））

③推理结果为0

进行代码规范，按照固定格式写会好些
https://blog.csdn.net/qq_14839543/article/details/114982354按照这种格式写
尤其是

with get_engine(engine_file) as engine, engine.create_execution_context() as context:

这句话部分，把下面的内容放在下面。

4.推理时batch_size的设置

①batch_size 为1时：

此时在./trtexec中不需要设置minshapes, optshapes, maxshapes等值也没问题，推理速度提升不少

②batch_size不为1时：

a.

在./trtexec中设置minshapes, optshapes, maxshapes的值https://zhuanlan.zhihu.com/p/446477075, 此时在推理时batch_size只要在min-max的范围内就可以，但是推理速度变慢，和纯gpu速度差不多了。

b.

在./trtexec中只设置opt_shapes，此时batch_size也就固定了，就是不为1了，此时在推理时batch_size好像就不能随意设置了

c.

解决推理时batch_size的问题：
不使用./trtexec命令行得到trt, 在py文件中生成trt， 此时设置min, opt, max, 推理速度不会降低

具体是：

# .pth转.onnx
def export_dynamic_onnx_model(model, inputs, dynamic_onnx_model_path):
    if os.path.exists(dynamic_onnx_model_path):
        return
    symbolic_names = {0: 'batch_size', 1: 'max_seq_len'} 
    with torch.no_grad():
        torch.onnx.export(model, 
                        (inputs['input_ids'], inputs['attention_mask'], inputs['token_type_ids']),
                        dynamic_onnx_model_path, 
                        verbose=True, do_constant_folding=True,
                        input_names=['input_ids', 'input_mask', 'input_token_type_ids'], 
                        output_names=['output'], opset_version=13,
                        dynamic_axes={'input_ids':symbolic_names, 'input_mask':symbolic_names, 
                        'input_token_type_ids':symbolic_names}) # , 'output':symbolic_names
    logger.info("ONNX Model exported to {0}".format(dynamic_onnx_model_path))




# onnx转trt
def onnx_build_engine(onnx_file_path, engine_file_path, args, write_engine = True):
    '''获取engine, 建立上下文'''
    G_LOGGER = trt.Logger(trt.Logger.WARNING)
    def build_engine():
        # 动态输入第一点必须要写的
        explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

        batch_size = 8
        with trt.Builder(G_LOGGER) as builder, builder.create_network(explicit_batch) as network, trt.OnnxParser(network, G_LOGGER) as parser:
            builder.max_batch_size = batch_size  # 推理最大能接受的batch_size
            config = builder.create_builder_config()
            config.max_workspace_size = common.GiB(2)
            config.set_flag(trt.BuilderFlag.FP16)
            print('Loading ONNX file from path {}...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:
                print('Beginning ONNX file parsing')
                parser.parse(model.read())
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))

            # 重点
            profile = builder.create_optimization_profile() # 动态输入时候需要 分别为最小输入、常规输入、最大输入
            # 注意自己有几个输入，有几个输入就要写几个profile.set_shpae, 名字和转onnx的时候相对应
            profile.set_shape('input_ids', (1,args.max_seq_length), (1,args.max_seq_length), (batch_size,args.max_seq_length))
            profile.set_shape('input_mask', (1,args.max_seq_length), (1,args.max_seq_length), (batch_size,args.max_seq_length))
            profile.set_shape('input_token_type_ids', (1,args.max_seq_length), (1,args.max_seq_length), (batch_size,args.max_seq_length))
            config.add_optimization_profile(profile)

            engine = builder.build_engine(network, config)
            print('Completed creating Engine')

            # 保存engine文件
            if write_engine:
                with open(engine_file_path, 'wb') as f:
                    f.write(engine.serialize())
            
        with open(engine_file_path, 'rb') as f, trt.Runtime(G_LOGGER) as runtime:
            engine = runtime.deserialize_cuda_engine(f.read())
            return engine
            
    if os.path.exists(engine_file_path):
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, 'rb') as f, trt.Runtime(G_LOGGER) as runtime:
            engine = runtime.deserialize_cuda_engine(f.read())
            return engine
    else:
        return build_engine()

5.推理

with onnx_build_engine(pred_config.onnx_model_path, pred_config.engine_model_path, args=args, write_engine=True) as engine, engine.create_execution_context() as context:
            context.active_optimization_profile = 0
            origin_inputshape=context.get_binding_shape(0)
            origin_inputshape[0],origin_inputshape[1]=inputs_ids.shape
            context.set_binding_shape(0, (origin_inputshape))
            context.set_binding_shape(1, (origin_inputshape))  # 这里报错是batchsize设置的不匹配，后面要设置成动态的
            context.set_binding_shape(2, (origin_inputshape)) # 有几个输入就写几个context.set_binding_shape

            '''输入数据填充'''
            inputs, outputs, bindings, stream = common.allocate_buffers_v2(engine, context)
            inputs[0].host = inputs_ids
            inputs[1].host = attention_mask
            inputs[2].host = token_type_ids

            '''tensorrt推理'''
            t1 = time.time()
            intent_logits, slot_logits = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
            t2 = time.time()
            trt_totaltime += (t2-t1)*1000
            print('trt_totaltime {} ms'.format(trt_totaltime))

tensorrt推理结果是list, 需要对list进行reshape，接下来就是对推理的结果与纯gpu时对比一下看效果了


fp32下快了不少。