FP32推理TensorRT演示

引言

YOLOv5最新版本的6.x已經(jīng)支持直接導(dǎo)出engine文件并部署到TensorRT上了。

FP32推理TensorRT演示

可能很多人不知道YOLOv5新版本6.x中已經(jīng)支持一鍵導(dǎo)出Tensor支持engine文件，而且只需要一條命令行就可以完成：演示如下：

python export.py --weights yolov5s.pt --include onnx engine --device 0

其中onnx表示導(dǎo)出onnx格式的模型文件，支持部署到：

- OpenCV DNN- OpenVINO- TensorRT- ONNXRUNTIME

但是在TensorRT上推理想要速度快，必須轉(zhuǎn)換為它自己的engine格式文件，參數(shù)engine就是這個(gè)作用。上面的命令行執(zhí)行完成之后，就會(huì)得到onnx格式模型文件與engine格式模型文件。--device 0參數(shù)表示GPU 0，因?yàn)槲抑挥幸粡埧?！上述?dǎo)出的FP32的engine文件。

使用tensorRT推理 YOLOv5 6.x中很簡(jiǎn)單，一條命令行搞定了，直接執(zhí)行：

python detect.py --weights yolov5s.engine --view-img --source data/images/zidane.jpg

FP16推理TensorRT演示

在上面的導(dǎo)出命令行中修改為如下

python export.py --weights yolov5s.onnx --include engine --half --device 0

其中就是把輸入的權(quán)重文件改成onnx格式，然后再添加一個(gè)新的參 --half 表示導(dǎo)出半精度的engine文件。就這樣直接執(zhí)行該命令行就可以導(dǎo)出生成了，圖示如下：

對(duì)比可以發(fā)現(xiàn)相比FP32大小的engine文件，F(xiàn)P16的engine文件比FP32的engine大小減少一半左右，整個(gè)文件只有17MB大小左右。

推理執(zhí)行的命令跟FP32的相同，直接運(yùn)行，顯示結(jié)果如下：

對(duì)比發(fā)現(xiàn)FP32跟FP16版本相比，速度提升了但是精度幾乎不受影響！

INT8量化與推理TensorRT演示

TensorRT的INT量化支持要稍微復(fù)雜那么一點(diǎn)點(diǎn)，最簡(jiǎn)單的就是訓(xùn)練后量化。只要完成Calibrator這個(gè)接口支持，我用的TensorRT版本是8.4.0.x的，它支持以下幾種Calibrator：

不同的量化策略，得到的結(jié)果可能稍有差異，另外高版本上的INT8量化之后到低版本的TensorRT機(jī)器上可能無(wú)法運(yùn)行，我就遇到過(guò)！所以建議不同平臺(tái)要統(tǒng)一TensorRT版本之后，再量化部署會(huì)比較好。上面的Calibrator都必須完成四個(gè)方法，分別是：

#使用calibrator驗(yàn)證時(shí)候每次張數(shù)，跟顯存有關(guān)系，最少1張get_batch_size  #獲取每個(gè)批次的圖像數(shù)據(jù)，組裝成CUDA內(nèi)存數(shù)據(jù)get_batch  #如果以前運(yùn)行過(guò)保存過(guò)，可以直接讀取量化，低碳給國(guó)家省電read_calibration_cache#保存calibration文件，量化時(shí)候會(huì)用到write_calibration_cache

這塊對(duì)函數(shù)集成不懂建議參考TensorRT自帶的例子：

TensorRT-8.4.0.6samplespythonint8_caffe_mnist

幾乎是可以直接用的！Copy過(guò)來(lái)改改就好了！

搞定了Calibrator之后，需要一個(gè)驗(yàn)證數(shù)據(jù)集，對(duì)YOLOv5來(lái)說(shuō)，其默認(rèn)coco128數(shù)據(jù)集就是一個(gè)很好的驗(yàn)證數(shù)據(jù)，在data文件夾下有一個(gè)coco128.yaml文件，最后一行就是就是數(shù)據(jù)集的下載URL，直接通過(guò)URL下載就好啦。

完成自定義YOLOv5的Calibrator之后，就可以直接讀取onnx模型文件，跟之前的官方轉(zhuǎn)換腳本非常相似了，直接在上面改改，最重要的配置與生成量化的代碼如下：

# build trt enginebuilder.max_batch_size = 1config.max_workspace_size = 1 << 30config.set_flag(trt.BuilderFlag.INT8)config.int8_calibrator = calibratorprint('Int8 mode enabled')plan = builder.build_serialized_network(network, config)

主要就是設(shè)置config中的flag為INT8，然后直接運(yùn)行，得到plan對(duì)象，反向序列化為engine文件，保存即可。最終得到的INT8量化engine文件的大小在9MB左右。

數(shù)據(jù)太少，只有128張， INT8量化之后的YOLOv5s模型推理結(jié)果并不盡如人意。但是我也懶得再去下載COCO數(shù)據(jù)集， COCO訓(xùn)練集一半數(shù)據(jù)作為驗(yàn)證完成的量化效果是非常好。 這里，我基于YOLOv5s模型自定義數(shù)據(jù)集訓(xùn)練飛鳥(niǎo)跟無(wú)人機(jī)，對(duì)得到模型，直接用訓(xùn)練集270張數(shù)據(jù)做完INT8量化之后的推理效果如下：

量化效果非常好，精度只有一點(diǎn)下降，但是速度比FP32的提升了1.5倍左右（3050Ti）。

已知問(wèn)題與解決

量化過(guò)程遇到這個(gè)錯(cuò)誤

[09/22/2022-2313] [TRT] [I]   Calibrated batch 127 in 0.30856 seconds.[09/22/2022-2316] [TRT] [E] 2: [quantization.cpp::70] Error Code 2: Internal Error (Assertion min_ <= max_ failed. )[09/22/2022-2316] [TRT] [E] 2: [builder.cpp::619] Error Code 2: Internal Error (Assertion engine != nullptr failed. )Failed to create the engineTraceback (most recent call last):

解決方法，把Calibrator中g(shù)etBtach方法里面的代碼：

img = np.ascontiguousarray(img, dtype=np.float32)

to

img = np.ascontiguousarray(img, dtype=np.float16)

這樣就可以避免量化失敗。