当前位置:   article > 正文

基于TensorRT和onnxruntime下pytorch的Bert模型加速对比实践_trtexc

trtexc

         写在前面,自从2021年4月19日,跳槽以来,一直在忙于公司的项目和业务,没怎么看论文玩模型了,也没有写博客了。最近一直都在搞elasticsearch相关的东西,以及模型加速方面的工作,觉得很有价值,我也获得了成长,同时试用期也顺利渡过,百忙之中,抽出时间,完成了下面这篇博客。

        由于业务需求上来了,需要对bert模型进行推理加速;现有的加速方案有很多,基于模型的压缩方向如模型蒸馏模型量化模型剪枝等,还有基于不修改模型结构(或者不需要自己动手修改)的前提下采用TensorRT或者onnxruntime来进行加速。在预研了onnxruntime和TensorRT的加速方法和做了很多实验的基础上,输出了一篇pytorch框架下bert模型加速对比效果,重点在于方案的尝试和实验结果的对比,以及代码的记录,原理上涉及的比较少。

一、onnxruntime和TensorRT简介

1、onnxruntime

        ONNXRuntime是微软推出的一款推理框架,用户可以非常便利的用其运行一个onnx模型,进行推理和训练。一般而言,先把其他的模型转化为onnx格式的模型,然后进行session构造,模型加载与初始化和运行。其推理时采用的数据格式是numpy格式,而不是tensor张量,当然onnxruntime可以才GPU上也可以在CPU上运行。

其加速原理采用官网上的一段话:

ONNX Runtime applies a number of graph optimizations on the model graph then partitions it into subgraphs based on available hardware-specific accelerators. Optimized computation kernels in core ONNX Runtime provide performance improvements and assigned subgraphs benefit from further acceleration from each Execution Provider.

主要是对模型图进行优化,同时基于特定的硬件加速器把模型图切分为更小的子图,使用onnxruntime核心进行计算算子的优化。

2、TensorRT

        TensorRT是英伟达公司针对自家的GPU产品开发的一个神经网络加速库。它只用于模型在GPU上的推理加速,不支持CPU,一般也不会用于模型的训练。TensorRT加速效果还是比较明显的,一般都是加速几倍、几十倍甚至上百倍。为何它能加速呢?

看上面官方给出一一张图:

1、Precision Calibration

精度校准——训练时由于梯度等对于计算精度要求较高,但是inference阶段可以利用精度较低的数据类型加速运算,降低模型的大小,例如FP16,int8,从而加速模型推理速度。

2、Layer & Tensor fusion

层和张量融合——TensorRT中将多个层的操作合并为同一个层,这样就可以一定程度的减少kernel launches和内存读写。比如把主流神经网络的conv、BN、Relu三个层融合为了一个层;把维度相同的张运算组合成另一个大的张量运算。每一层的运算操作都是由GPU完成的——GPU通过启动不同的CUDA(Compute unified device architecture)核心来完成计算的,CUDA核心计算张量的速度是很快的,但是往往大量的时间是浪费在CUDA核心的启动和对每一层输入/输出张量的读写操作上面,这造成了内存带宽的瓶颈和GPU资源的浪费。

3、Kernel Auto-Tuning

计算核心自动调整——TensorRT可以针对不同的算法,不同的网络模型,不同的GPU平台,进行 CUDA核的调整,以保证当前模型在特定平台上以最优性能计算。

4、 Dynamic Tensor Memory

动态张量显存——每个tensor的使用期间,TensorRT会为其指定显存,避免显存重复申请,减少内存占用和提高重复使用效率。

5、Multi-Stream Execution

并行处理多流输入——这个就是GPU底层优化,理解不了。

6、 Time Fusion

时间融合——使用动态生成的算子优化循环神经网络。

以上对为何能加速进行了简单的介绍,详细的原理很难有比较深刻的理解。总体就是量化——降低数据精度、cuda kernel 智能化计算、动态显存管理以及模型结构和张量融合突破GPU带宽瓶颈。

二、pytorch+onnxruntime的bert模型加速效果

onnxruntime要想加速pytorch下的bert模型,首先就需要把pytorch下的bert模型转化为.onnx模型文件,torch自带的框架就能完成这个工作;接下来就是构建session,把转化好的.onnx文件加载进来,喂入数据进行推理。

环境:linux centos、cuda11.2、pytorch1.8、python3.7、3090显卡、onnxruntime-gpu 1.8.1

直接上代码:

torch转化为onnx,并构建InferenceSession推理:

  1. MODEL_ONNX_PATH = "../../onnx/torch_bert_base_fixed_" + str(batch_size) + ".onnx"
  2. torch.onnx.export(model, org_dummy_input, MODEL_ONNX_PATH, verbose=True,
  3. input_names=['input_ids', 'attention_mask'],
  4. output_names=['output'], opset_version=11)
  5. print("Export of torch_model.onnx complete!")
  6. onnx_session = onnxruntime.InferenceSession(MODEL_ONNX_PATH)
  7. pred_onnx = onnx_session.run(None, {'input_ids': inf_dummy_input[0], 'attention_mask': inf_dummy_input[1]})

完整代码如下:

  1. import torch
  2. import onnxruntime
  3. from transformers import BertModel
  4. import time
  5. import matplotlib.pyplot as plt
  6. from torch.quantization import quantize_dynamic
  7. import sys
  8. sys.path.append("../../")
  9. from codes.dataReader.entity_data_reader import EntityDataReader
  10. from torch.utils.data import DataLoader
  11. import argparse
  12. from tqdm import tqdm
  13. def make_train_dummy_input(dataloader,device):
  14. for batch in dataloader:
  15. batch = [t.to(device) for t in batch]
  16. dummy_input_ids = batch[0]
  17. dummy_attention_masks = batch[1]
  18. del batch
  19. break
  20. return (dummy_input_ids,dummy_attention_masks)
  21. def make_inference_dummy_input(dataloader,device):
  22. for batch in dataloader:
  23. batch = [t.to(device) for t in batch]
  24. inf_input_ids = batch[0]
  25. inf_attention_masks = batch[1]
  26. del batch
  27. break
  28. return (inf_input_ids, inf_attention_masks)
  29. def set_args():
  30. parser = argparse.ArgumentParser(description='train_bert_entity_classification args')
  31. parser.add_argument('--model_path',type=str,default='../../pretrain_models/chinese-bert-wwm-ext')
  32. parser.add_argument('--batch_size',default=128 ,type = int)
  33. parser.add_argument('--test_file_path', default='../../data/data_clean/entity_train.xlsx', type=str)
  34. args = parser.parse_args()
  35. return args
  36. if __name__ == '__main__':
  37. args = set_args()
  38. device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
  39. count = 1001
  40. model = BertModel.from_pretrained(args.model_path)
  41. model.to(device)
  42. model.train(False)
  43. test_dataset = EntityDataReader(args.test_file_path, args.model_path)
  44. org_dummy_input = make_train_dummy_input(test_dataloader, device)
  45. org_meantimes = []
  46. onnx_meantimes = []
  47. batch_sizes = [ 2**i for i in range(0,9)]
  48. for batch_size in batch_sizes:
  49. test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
  50. #org_dummy_input = make_train_dummy_input(test_dataloader, device)
  51. inf_dummy_input = make_inference_dummy_input(test_dataloader, device)
  52. MODEL_ONNX_PATH = "../../onnx/torch_bert_base_fixed_" + str(batch_size) + ".onnx"
  53. torch.onnx.export(model, org_dummy_input, MODEL_ONNX_PATH, verbose=True,
  54. input_names=['input_ids', 'attention_mask'],
  55. output_names=['output'], opset_version=11)
  56. print("Export of torch_model.onnx complete!")
  57. onnx_session = onnxruntime.InferenceSession(MODEL_ONNX_PATH)
  58. totaltime = 0
  59. for i in tqdm(range(count), ncols=50):
  60. t1 = time.time()
  61. output = model(*inf_dummy_input)[0]
  62. if i ==0:
  63. print(output.shape)
  64. del output
  65. t2 = time.time()
  66. if i >0:
  67. totaltime += (t2 - t1)
  68. org_meantime = totaltime / count
  69. print('Bert inference mean time is %.4f' % (totaltime / count))
  70. inf_dummy_input = [t.cpu().numpy() for t in inf_dummy_input]
  71. totaltime = 0
  72. for i in tqdm(range(count), ncols=50):
  73. t1 = time.time()
  74. pred_onnx = onnx_session.run(None, {'input_ids': inf_dummy_input[0], 'attention_mask': inf_dummy_input[1]})
  75. t2 = time.time()
  76. if i == 0:
  77. print(pred_onnx[0].shape)
  78. del pred_onnx
  79. if i >0 :
  80. totaltime += (t2 - t1)
  81. onnx_meantime = totaltime / (count-1)
  82. print('Bert inference by onnxruntime mean time is %.4f' % (totaltime / (count-1)))
  83. onnx_meantimes.append(onnx_meantime*1000)
  84. org_meantimes.append(org_meantime*1000)
  85. del org_dummy_input
  86. del inf_dummy_input
  87. del onnx_session
  88. torch.cuda.empty_cache()
  89. plt.plot(batch_sizes,org_meantimes,color='red',label='bert_base infr time')
  90. plt.plot(batch_sizes, onnx_meantimes, color='blue', label='bert_base onnx infr time')
  91. plt.grid(alpha=0.4, linestyle=':')
  92. plt.legend(loc="upper right")
  93. plt.ylabel("inference time/ms")
  94. plt.xlabel("batch size")
  95. plt.savefig('../../onnx/acceleration_comparsion.png')
  96. plt.show()

结果如下图:

seq_length=100,hidden_state=768 ,可以看到加速效果只有在batch-size小于10以及32-160之间

才有效果,而且是batch_size = 1 的时候比较明显。

seq_length=100,hidden_state=256的时候,结果不一样,如下图(bert-small版本):

图上的时间单位应该是ms,batch_size = 1的时候加速效果大概是5倍,<64batch_size<256两者差不多,batch_size>256以后3090上纯GPU推理更快。

纯GPU下,可以得出batch_size<某个值的时候,推理时间一样;后面随着batch_size增大线性增加,当然也和hidden_state密切相关。

onnxruntime在小batch_size下还是有一定的加速效果的,相对纯GPU而言。

三、pytorch+onnx+TensorRT的bert模型加速效果

这个技术方案稍微麻烦一点,需要把pytorch转成的onnx模型再转化为tensorrt下的.trt文件,这里又又两种方法(其实是一种,一种需要代码,一种不需要代码);

1、采用TensorRT自带的trtexc把onnx转化为.trt

2、调用tensorrt的API把onnx转化为.trt

以上方法在转化过程中,bert构建engine的时候,对tensorRt的版本要求比较高,我这里使用的是8.0.1.6,其他的低版本都试过,转化的过程中都会报错,不支持pytorch中bert某些算子,需要修改bert的源码,也是比较麻烦。

当然还有一种更直接的办法就是使用tensorrt的API直接把pytorch下的bert模型构建成一个trt的engine执行推理,这个需要对tensorrt的API非常熟悉,以及bert模型的权重和结构非常熟悉,对代码功底要求也比较高,这个等以后有时间了再来实现,简单的图片分类模型例子可以参考TensorRt 官方github给出的例子——目前没有时间,后面自己一定要实现一遍,对自己很有作用。

环境:linux centos、cuda11.2、pytorch1.8、python3.7、3090显卡、tensorrt8.0.1.6

1、采用TensorRT自带的trtexc把onnx转化为.trt

在linux系统安装好的TensorRT-8.0.1.6,bin目录中,可以看到trtexec:

 转化命令:

./trtexec --onnx=onnxModelFilePath --saveEngine=trtEngineFliePath

这样就能把一个.onnx模型转化为tensorRT的engine文件

给出一个示例:

./trtexec --onnx=/home/AI_team/yanghuang/competition/onnx/torch_bert_base_fixed_256.onnx --maxBatch=200 --workspace=1000 --fp16 --saveEngine=/home/AI_team/yanghuang/competition/tensorrt_engine/torch_bert_base_fixed_256.trt

--onnx:指定onnx模型路径

--maxBatch=200:指定trt最大的batch_size=200

--workspace=1000:指定转化过程中的工作空间是1000M

--fp16:指定采用了fp16精度——半精度,也还可以是int8

--saveEngine:指定trt文件保存的路径

以上就是常用的参数,其中--onnx和--saveEngine是必须的,其他可选(因为有默认选项)

还有很多其他的参数,它们具体是什么有什么作用,需要读者自己去研究了./trtexec --h就可以查看全部的参数。

对于很大的模型不好开启详细的转化日志,不然就要等很久了

转化过程和结果截图:

转化过程中,GPU参与运算

可以看到batch_size = 256的时候,bert_base版本模型推理吞吐量Throughout=12.22,推理平均时间是:86毫秒,也就是截图中的Latency——延迟——GPU计算时间+H2D+D2H处理一个query的时间之和。

2、调用tensorrt的API把onnx转化为.trt

主要的API是

trt.Builder(TRT_LOGGER) as builder创建一个builder
builder.create_network(explicit_batch) as network创建一个空的网络
trt.OnnxParser(network, TRT_LOGGER) as parser创建一个onnx解析器
trt.Runtime(TRT_LOGGER) as runtime创建一个trt的运行环境
config = builder.create_builder_config()创建builder_config

设置一些属性后,就可以直接解析onnx然后转化为trt engine

plan = builder.build_serialized_network(network, config)
engine = runtime.deserialize_cuda_engine(plan)

序列化和反序列化engine

完整代码

  1. def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="", fp16_mode=True, int8_mode=False, save_engine=False):
  2. """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
  3. explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)#这句一定需要加上不然会报错
  4. def build_engine(max_batch_size, save_engine):
  5. """Takes an ONNX file and creates a TensorRT engine to run inference with"""
  6. with trt.Builder(TRT_LOGGER) as builder, \
  7. builder.create_network(explicit_batch) as network, \
  8. trt.OnnxParser(network, TRT_LOGGER) as parser, \
  9. trt.Runtime(TRT_LOGGER) as runtime:
  10. print(network.num_layers)
  11. print(network.num_inputs)
  12. print(network.num_outputs)
  13. print(network.name)
  14. config = builder.create_builder_config()
  15. config.max_workspace_size = 1 << 30 # 256MiB
  16. builder.max_batch_size = max_batch_size
  17. if fp16_mode:
  18. config.set_flag(trt.BuilderFlag.FP16)
  19. elif int8_mode:
  20. config.set_flag(trt.BuilderFlag.INT8)
  21. else:
  22. config.set_flag(trt.BuilderFlag.REFIT)
  23. flag = builder.is_network_supported(network,config)
  24. print('flag',flag)
  25. # Parse model file
  26. if not os.path.exists(onnx_file_path):
  27. quit('ONNX file {} not found'.format(onnx_file_path))
  28. print('Loading ONNX file from path {}...'.format(onnx_file_path))
  29. with open(onnx_file_path, 'rb') as model:
  30. print('Beginning ONNX file parsing')
  31. # print(type(model.read()))
  32. parser.parse(model.read())
  33. # parser.parse_from_file(onnx_file_path)
  34. assert network.num_layers > 0, 'Failed to parse ONNX model.Please check if the ONNX model is compatible '
  35. # last_layer = network.get_layer(network.num_layers - 1)
  36. # network.mark_output(last_layer.get_output(0))
  37. print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
  38. plan = builder.build_serialized_network(network, config)
  39. engine = runtime.deserialize_cuda_engine(plan)
  40. print("Completed creating Engine")
  41. if save_engine:
  42. with open(engine_file_path, "wb") as f:
  43. f.write(plan)
  44. return engine
  45. if os.path.exists(engine_file_path):
  46. # If a serialized engine exists, load it instead of building a new one.
  47. print("Reading engine from file {}".format(engine_file_path))
  48. with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
  49. return runtime.deserialize_cuda_engine(f.read())
  50. else:
  51. return build_engine(max_batch_size, save_engine)

注意——

explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

这句代码一定需要指定,不然会报错。

3、tensorrt进行推理加速

使用tensorrt进行推理加速稍微有点麻烦,总结起来步骤就是:

a、获取engine,建立上下文

  1. engine = get_engine(engine_model_path)
  2. context = engine.create_execution_context()

b、从engine中获取inputs, outputs, bindings, stream的格式以及分配缓存

  1. inputs, outputs, bindings, stream = common.allocate_buffers(engine)
  2. class HostDeviceMem(object):
  3. def __init__(self, host_mem, device_mem):
  4. self.host = host_mem
  5. self.device = device_mem
  6. def __str__(self):
  7. return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
  8. def __repr__(self):
  9. return self.__str__()
  10. # Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
  11. def allocate_buffers(engine):
  12. inputs = []
  13. outputs = []
  14. bindings = []
  15. stream = cuda.Stream()
  16. for binding in engine:
  17. size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
  18. dtype = trt.nptype(engine.get_binding_dtype(binding))
  19. # Allocate host and device buffers
  20. host_mem = cuda.pagelocked_empty(size, dtype)
  21. device_mem = cuda.mem_alloc(host_mem.nbytes)
  22. # Append the device buffer to device bindings.
  23. bindings.append(int(device_mem))
  24. # Append to the appropriate list.
  25. if engine.binding_is_input(binding):
  26. inputs.append(HostDeviceMem(host_mem, device_mem))
  27. else:
  28. outputs.append(HostDeviceMem(host_mem, device_mem))
  29. return inputs, outputs, bindings, stream

c、输入数据填充

在第二步从engine得到了输入输出的数据格式以及分配了缓存,要使用其他的推理数据,就要把新的数据填充进去

  1. inputs, outputs, bindings, stream = common.allocate_buffers(engine)
  2. inf_dummy_inputs_all = []
  3. for i in tqdm(range(count),ncols=50):
  4. inf_dummy_inputs = make_inference_dummy_input(test_dataloader,device)
  5. inf_dummy_inputs_all.append(inf_dummy_inputs)
  6. for input,inf_dummy_input in zip(inputs,inf_dummy_inputs):
  7. temp = inf_dummy_input.view(-1).numpy().astype(np.int32)
  8. input.host = temp

注意这里输入的数据格式是np的int32,不是张量也不是int64。

d、tensorrt推理

细节在代码中有详细的注释

  1. # This function is generalized for multiple inputs/outputs.
  2. # inputs and outputs are expected to be lists of HostDeviceMem objects.
  3. def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
  4. # Transfer input data to the GPU.
  5. [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
  6. # Run inference.
  7. context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
  8. # Transfer predictions back from the GPU.
  9. [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
  10. # Synchronize the stream
  11. stream.synchronize()
  12. # Return only the host outputs.
  13. return [out.host for out in outputs]

完整的tensorrt推理和纯GPU效果对比代码:

  1. import torch
  2. from transformers import BertModel
  3. import time
  4. import logging # 引入logging模块
  5. logging.basicConfig(level=logging.DEBUG,format='%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s')
  6. import sys
  7. sys.path.append('../../')
  8. logging.info('{}'.format(sys.path))
  9. from codes.dataReader.entity_data_reader import EntityDataReader
  10. import codes.onnx_trt.common as common
  11. from torch.utils.data import DataLoader
  12. import argparse
  13. from tqdm import tqdm
  14. import tensorrt as trt
  15. TRT_LOGGER = trt.Logger() # This logger is required to build an engine
  16. import random
  17. import numpy as np
  18. def make_train_dummy_input(dataloader,device):
  19. for batch in dataloader:
  20. batch = [t.to(device) for t in batch]
  21. dummy_input_ids = batch[0]
  22. dummy_attention_masks = batch[1]
  23. break
  24. return (dummy_input_ids,dummy_attention_masks)
  25. def make_inference_dummy_input(dataloader,device):
  26. for batch in dataloader:
  27. batch = [t.to(device) for t in batch]
  28. inf_input_ids = batch[0]
  29. inf_attention_masks = batch[1]
  30. break
  31. return (inf_input_ids, inf_attention_masks)
  32. def set_args():
  33. parser = argparse.ArgumentParser(description='train_bert_entity_classification args')
  34. parser.add_argument('--model_path',type=str,default='../../pretrain_models/chinese-bert-wwm-ext')
  35. parser.add_argument('--batch_size',default=256 ,type = int)
  36. parser.add_argument('--test_file_path', default='../../data/data_clean/entity_train.xlsx', type=str)
  37. args = parser.parse_args()
  38. return args
  39. def get_engine(engine_file_path):
  40. print("Reading engine from file {}".format(engine_file_path))
  41. with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
  42. return runtime.deserialize_cuda_engine(f.read())
  43. def create_random_sample(batch_size):
  44. batch = []
  45. for i in range(batch_size):
  46. t = [2]
  47. for j in range(98):
  48. t.append(random.randint(0, 14467))
  49. t.append(3)
  50. batch.append(t)
  51. input_ids = torch.tensor(batch, dtype=torch.long)
  52. token_type_ids = torch.zeros((batch_size, 100), dtype=torch.long)
  53. attention_mask = torch.ones((batch_size, 100), dtype=torch.long)
  54. return (input_ids, token_type_ids, attention_mask)
  55. if __name__ == '__main__':
  56. args = set_args()
  57. test_dataset = EntityDataReader(args.test_file_path, args.model_path)
  58. test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)
  59. device = 'cpu'
  60. count = 10
  61. trt_results = []
  62. nue_results = []
  63. # batch_size = 100
  64. nue_gpu_totaltime = 0
  65. nue_trt_totaltime = 0
  66. engine_model_path = "../../tensorrt_engine/torch_bert_base_fixed_256.trt"
  67. engine = get_engine(engine_model_path)
  68. context = engine.create_execution_context()
  69. inputs, outputs, bindings, stream = common.allocate_buffers(engine)
  70. inf_dummy_inputs_all = []
  71. for i in tqdm(range(count),ncols=50):
  72. inf_dummy_inputs = make_inference_dummy_input(test_dataloader,device)
  73. inf_dummy_inputs_all.append(inf_dummy_inputs)
  74. for input,inf_dummy_input in zip(inputs,inf_dummy_inputs):
  75. temp = inf_dummy_input.view(-1).numpy().astype(np.int32)
  76. input.host = temp
  77. t1 = time.time()
  78. output = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
  79. t2 = time.time()
  80. nue_trt_totaltime += (t2-t1)
  81. output[0] = output[0].reshape(args.batch_size, -1)
  82. if i==0:
  83. print('output[0].shape',output[0].shape)
  84. del output
  85. torch.cuda.empty_cache()
  86. device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
  87. nue_model = BertModel.from_pretrained(args.model_path)
  88. nue_model.to(device)
  89. nue_model.eval()
  90. with torch.no_grad():
  91. for i in tqdm(range(count),ncols=50):
  92. inf_dummy_inputs = inf_dummy_inputs_all[i]
  93. inf_dummy_inputs = [t.to(device) for t in inf_dummy_inputs]
  94. t1 = time.time()
  95. logits = nue_model(*inf_dummy_inputs)[0]
  96. nue_results.append(logits.detach().cpu().numpy())
  97. t4 = time.time()
  98. if i==0:
  99. print('logits.shape',logits.shape)
  100. nue_gpu_totaltime += (t4 - t1)
  101. del inf_dummy_inputs
  102. del logits
  103. torch.cuda.empty_cache()
  104. print("bert inference by 3090GPU with tensorrt per batch time is %.4f"%(nue_trt_totaltime/count))
  105. print("bert inference by 3090GPU per batch time is %.4f" % (nue_gpu_totaltime / count))

上面的代码有一个值得注意的点就是——

common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)

推理的代码do_inference运行之前不要设置device为GPU上的cuda,不然就会出现错误

[TensorRT] ERROR: 1: [slice.cu::launchNaiveSliceImpl::148] Error Code 1: Cuda Runtime (invalid resource handle)

 最终对比结果如下:

a、采用./trtexec得到的trt engine 

engine_model_path = "../../tensorrt_engine/torch_bert_base_fixed_256.trt"

 采用tensorrt使用3090来推理batch_size=256,fp16用时85.8ms;纯GPU模式下240.6ms,提升3倍,这里和./trtexec转化onnx为trt过程中给出的推理时间是一致的

b、采用tensorRT API得到的trt engine

engine_model_path = "../../tensorrt_engine/torch_bert_base_fixed_256_api.trt"

 采用tensorrt使用3090来推理batch_size=256,fp16用时86.1ms;纯GPU模式下240.3ms,提升3倍,这里和./trtexec转化onnx为trt过程中给出的推理时间也是一致的,和上面的a方案也是一致的,本次实验应该是正确无误的。

加速效果还是比较明显的,当然这里可能有一点点精度损失,在不采用fp16的情况下,我们的业务模型的精度是99.72%。

以上就是这段时间预研tensorrt的收获,算是能走通流程,能搭建一定的任务了。

关于int8模式没有去研究,套路应该是一样的,不过比较麻烦的应该是需要一个精度校准器Calibrator,后面有需要或者有空的时候可以研究一下,应该不会花太多时间。

参考文章

TensorRT-优化-原理

PyTorch模型转TensorRT是怎么实现的?

onnxruntime官网

英伟达TensorRT官网

tensorrt的API文档

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/繁依Fanyi0/article/detail/104712
推荐阅读
相关标签
  

闽ICP备14008679号