如何将自己训练的yolo模型转为tensorrt相关的模型
本文主要是以自己训练的yolov3-tiny模型作为例子进行总结和方法说明,其中方法也是照搬了github上面别人的总结,在这里写主要是为了写一些自己认为值得注意的点以及自己在转换过程中遇到的问题,同时也方便自己之后的使用,不然又得忘光光,果然好记性不如烂笔头啊!!
【1】.pt模型转.wts模型
注意:在实现.pt模型转.wts模型的过程中,首先要考虑的问题是自己训练出来的模型存储的时候是只保存了权重参数的.pt模型还是权重参数和结构都保存的.pt模型
①对于只保存了权重参数的.pt模型
如果只保存了权重参数,这个时候可以结合.cfg文件一起实现模型的转换,转换的方法如下:
1)从github里clone一个yolov3.git并解压
git clone -b archive https://github.com/ultralytics/yolov3.git
2)从github里clone一个tensorrtx.git并解压
git clone https://github.com/wang-xinyu/tensorrtx.git
3)将自己训练完的yolov3-tiny模型copy到解压后的yolov3/weights文件夹中,同时将tensorrtx/yolov3-tiny文件夹里面的gen_wts.py文件copy到yolov3的文件夹中
4)在yolov3的文件夹中运行gen_wts.py即可将.pt模型转换成.wts模型
python3 gen_wts.py ./weights/XXX.pt