yolov5的pt权重转tensorrt的trt权重

yolov5的pt权重转tensorrt的trt权重

相信如何利用tensorrt进行加速会是大家提高网络速度的关键一环，实际步骤其实也只需 pt/pth 转到 onnx ，再onnx转trt即可。虽然只有两步，但确实会挺折磨人，在这我就浅浅分享一个大致的步骤，以yolov5为例，希望能给大家一点点启发。

pt转onnx

这一环，yolov5中自带的export.py可以很出色的帮助我们完成转换，只需要执行:

python export.py --weights E:\Torch2Trt\test.pt --img 640 --batch 1

就可以在pt文件同路径下生成onnx文件，如果希望迁移到大家自身的模型，个人觉得可以参考一下export.py中的代码。

onnx转trt

这一步需要大家提前把tensorrt下载并且编译好，这里我就不做过多描述大家可以参考网上他人写的安装教程，我这里下的是TensorRT-8.2.4.2，编译起来挺方便的。安装完成后大家需要切换到TensorRT-8.2.4.2下面的bin文件夹：

可以使用下述代码进行转换，具体的参数其实还挺多的，需要大家自己去了解了再使用。

trtexec --onnx=E:\Torch2Trt\TensorRT-8.2.4.2\bin\test.onnx --explicitBatch --saveEngine=E:\Torch2Trt\best.trt --workspace=1024 --best

但直接跑上面这段代码是走不通的，它会报错下面这段错误:Assertion failed: scales.is_weights() && “Resize scales must be an initializer!”

因为默认的yolov5的模型中upsample不是指定倍率进行放大的，我们需要修改models文件夹下的tf.py文件找到下述这段代码：

注释掉红色箭头所指的代码，消掉黄色箭头所指代码的注释符，然后保持重新训练得到新的模型再走上述的步骤就可以走得通了。大家在转换其它模型的时候，如果遇到类似的问题，但不是由Upsample引起的，可以参考这篇知乎文章
上面还总结了其它类似的错误。