Melody-Zhou / tensorRT_Pro-YOLOv8

This repository is based on shouxieai/tensorRT_Pro, with adjustments to support YOLOv8.
MIT License
240 stars 40 forks source link

tensorrt多batch推断 #26

Closed BaseJiang closed 5 months ago

BaseJiang commented 5 months ago

作者你好,在转onnx的时候设置了动态的batch,在onnx编译为tensorrt模型的时候,有个参数"test_batch_size",个人理解该参数是固定最大的batch推断数。在将其设置为2时候,编译的模型可以正常同时推断两张图,也可以正常推断一张图。但是test_batch_size设置为1,编译出的模型,在测试的时候,发现也可以同时推断多张图像,请问这是什么原因,还是说在onnx模型时已经固定了多batch推断,转trt模型的时候,不需要关注这个参数?

Melody-Zhou commented 5 months ago

作者你好,在转onnx的时候设置了动态的batch,在onnx编译为tensorrt模型的时候,有个参数"test_batch_size",个人理解该参数是固定最大的batch推断数。在将其设置为2时候,编译的模型可以正常同时推断两张图,也可以正常推断一张图。但是test_batch_size设置为1,编译出的模型,在测试的时候,发现也可以同时推断多张图像,请问这是什么原因,还是说在onnx模型时已经固定了多batch推断,转trt模型的时候,不需要关注这个参数?

compile 函数中,maxBatchSize 参数主要用于以下两个方面:

  1. 设置 build 的最大 batch

我们在利用 trtexec 构建 engine 时如果是动态 onnx 模型则需要设置 --minShapes--optShapes--maxShapes,这里的 maxBatchSize 参数会被用来设置 --maxShapes

如果 maxBatchSize 设置为 16,那么在推理时可以一次性处理最多 16 张图片:

builder->setMaxBatchSize(maxBatchSize);
  1. 校准器的 batch

在 INT8 模式下需要进行校准以确定量化参数,而校准器需要知道 batch 的大小以便正确处理输入数据。因此,在构造 INT8 校准器(Int8EntropyCalibratorInt8MinMaxCalibrator)时,需要使用 maxBatchSize 来设置校准输入的维度。不同的 batch size 可能会影响 PTQ 量化的精度

auto calibratorDims = inputDims;
calibratorDims.d[0] = maxBatchSize;

另外 test_batch_size 设置为 1 时也可以同时推断多张图像是因为这些图片应该是一张张塞进去推理的,而不是一个 batch 塞进去的,你看下相关耗时就知道了,我在 NX 上测试自训练的一个模型在 batch=1 和 batch=16 的耗时如下所示:

Jetson Xavier NX
FP16 batch16 3.41ms
FP16 batch1 5.04ms
BaseJiang commented 5 months ago

感谢解答,明天测试下不同batch的推断时间统计下。