在前面两篇我们介绍了TensorRT的环境部署以及TensorRT的功能性介绍。
在使用tensorRT的时候,最常见的方式就是各种训练框架都基于Onnx来做中间转换,通过Onnx来生成TensorRT engine,进而享受到TensorRT的推理速度。所以本文就以常见的检测模型yolov5 6.0 来进行TensorRT模型的部署。整体测试都是基于TensorRT python backend,在开始之间请大家思考如下问题?
1.推理时所选的BatchSize 和 Fps的关系是怎样的?
2.TensorRT推理一定比pytorch推理快吗,我们要无脑替换成TensorRT是正确的使用姿势吗?受到哪些因素的影响?对于使用python作为主要开发语言的同学来说,什么时候考虑使用tensorRT作为推理引擎,什么时候使用pytorch呢?
带着这些问题,我们共同开始今天的学习。本文的实验环境是GPU V100S 32G,Intel(R) Xeon(R) Gold 6248R CPU @ 3.00GHz 。本文选用的方式是Explict batchSize,batchsize是从onnx模型继承过来的,关于Explicit 和 Implicit的方式,我们后续再进行介绍。
关于onnx模型导出的问题,我觉得还是另开篇章去讲述吧。export的代码yolov5官方已经提供了。大家可以参考一下这份代码:
yolov5官方onnx的转换代码
使用方法:
python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite
重要参数:
我们需要测试不同batchsize下面的tensorRT模型,所以导出的时候指定不同的batchsize。(注意这是在Explicit batch模式下进行的)。另外我这边选择测试的是fp16,所以需要都加上--half参数。
最简单的方法就是使用trtexec作为onnx转换的工具。trtexec是Nvidia TensorRT自带的命令行工具。一方面可以帮助我们进行onnx或者其他格式的模型快速转换到TensorRT engine,包括指定batchsize、precision等,另一方面可以帮助我们测试模型的基准性能。
因为我要测试不同batchsize下,tensorRT与pytorch之间的推理性能差距,所以在导出的时候需要使用对应的onnx模型。
trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s_engine_fp16_b64.trt --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --explicitBatch --fp16
这里因为是测试半精度,所以需要--inputIOFormats=fp16:chw --outputIOFormats=fp16:chw,这两个参数。我这里遇到一个坑就是,只指定--fp16,另外两个参数如果不指定的话,得到的模型推理速度会很慢。查阅文档是说默认是按照fp32的精度进行,即使单独指定了--fp16参数。
我分别测试了4、8、16、32、64五组batchsize,得到了对应的5个engine模型。
另一种方式是基于tensorRT的python接口进行engine转换。
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def ONNX_to_TRT(onnx_model_path=None, trt_engine_path=None, fp16_mode=False):
"""
仅适用TensorRT V8版本
生成cudaEngine,并保存引擎文件(仅支持固定输入尺度)
Serialized engines不是跨平台和tensorRT版本的
fp16_mode: True则fp16预测
onnx_model_path: 将加载的onnx权重路径
trt_engine_path: trt引擎文件保存路径
"""
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(EXPLICIT_BATCH)
parser = trt.OnnxParser(network, TRT_LOGGER)
config = builder.create_builder_config()
config.max_workspace_size = GiB(1)
if fp16_mode:
config.set_flag(trt.BuilderFlag.FP16)
with open(onnx_model_path, 'rb') as model:
assert parser.parse(model.read())
serialized_engine = builder.build_serialized_network(network, config)
with open(trt_engine_path, 'wb') as f:
f.write(serialized_engine) # 序列化
print('TensorRT file in ' + trt_engine_path)
print('============ONNX->TensorRT SUCCESS============')
class TrtModel():
'''
TensorRT infer
'''
def __init__(self, trt_path, device=1):
self.cfx = cuda.Device(device).make_context()
self.stream = cuda.Stream()
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
# 启动一个tensorRT pyton runtime
runtime = trt.Runtime(TRT_LOGGER)
# 反序列化模型,使用runtime加载模型
with open(trt_path, "rb") as f:
self.engine = runtime.deserialize_cuda_engine(f.read())
# 创建执行上下文
self.context = self.engine.create_execution_context()
# 模型输入的尺寸
intype = self.engine.get_binding_dtype("input")
insize = trt.volume(self.engine.get_binding_shape("input"))
# fp16 占2个字节 fp32 占4个字节
insize = insize * 2 if intype == DataType.HALF else insize * 4
# 模型输出的尺寸
otype = self.engine.get_binding_dtype("output")
osize = trt.volume(self.engine.get_binding_shape("output"))
osize = osize * 2 if otype == DataType.HALF else osize * 4
otype = np.float16 if otype == DataType.HALF else np.float32
# 分配输入输出的显存
self.cuda_mem_input = cuda.mem_alloc(insize)
self.cuda_mem_output = cuda.mem_alloc(osize)
self.bindings = [int(self.cuda_mem_input), int(self.cuda_mem_output)]
self.output = np.empty(self.engine.get_binding_shape("output"), dtype=otype)
def __call__(self, img_np_nchw):
'''
TensorRT推理
:param img_np_nchw: 输入图像
'''
self.cfx.push()
#将数据从内存拷贝到显存中
cuda.memcpy_htod_async(self.cuda_mem_input, img_np_nchw.ravel(), self.stream)
#tensorRT异步并发推理
self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
#将数据从显存传输到内存
cuda.memcpy_dtoh_async(self.output, self.cuda_mem_output, self.stream)
# 等待所有cuda核完成计算
self.stream.synchronize()
self.cfx.pop()
return self.output
def destroy(self):
# Remove any context from the top of the context stack, deactivating it.
self.cfx.pop()
这是一个简化版本的tensorRT推理类。其原理其实就是提前开辟显存空间,然后把数据从内存拷贝到显存。在显卡上面利用cuda核进行并行计算,然后在从显存把结果数据拷贝到内存。
另外关于pagelock memory(锁页内存)的使用,有些代码里会提前开辟锁页内存,我在实际测试的时候发现,并没有特别的速度提升。
我在实际测试的时候发现,对于不同的batchsize,在tensorRT的python backend上面,不是完全都呈现出比pytorch推理更快的现象,而是呈现出有规律的趋势。在batchsize比较小的时候,tensorRT推理具备碾压的性能优势,但是随着batchsize的增大,这个gap在慢慢变小,在大于32、64等的情况下,反而没有pytorch直接推理快。
在c++上面的测试,整体上tensorRT的推理都比pytorch快,但是同样这个gap在减少。
有知道原因的朋友,求点拨。下面是我记录的实验记录。
可以明显的看到这个趋势。所以在用python封装backend的时候,感觉tensorRT更适合小batchsize的使用场景,比如单张图片的某种推理服务。如果是视频推理的情景,可能不间断要一批批帧画面进行推理,对于这种情况,用pytorch的半精度直接推理可能效果反而更好,并不是无脑使用的。
以上是以yolov5 6.0做的实验,来记录python backend的使用方法和记录测试结果。谢谢阅读。