2022-10-31如何将自己训练的yolo模型转为tensorrt相关的模型（.trt or .engine）

本文主要是以自己训练的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

此处附上gen_wts.py的内容

import struct
import sys
from models import *
from utils.utils import *

model = Darknet('cfg/yolov3-tiny.cfg', (608, 608))
weights = sys.argv[1]
device = torch_utils.select_device('0')
if weights.endswith('.pt'):  # pytorch format
    model.load_state_dict(torch.load(weights, map_location=device)['model'])
else:  # darknet format
    load_darknet_weights(model, weights)
model = model.eval()

with open('yolov3-tiny.wts', 'w') as f:
    f.write('{}\n'.format(len(model.state_dict().keys())))
    for k, v in model.state_dict().items():
        vr = v.reshape(-1).cpu().numpy()
        f.write('{} {} '.format(k, len(vr)))
        for vv in vr:
            f.write(' ')
            f.write(struct.pack('>f',float(vv)).hex())
        f.write('\n')

注意：在转换的过程中.cfg文件要换成自己训练的对应的.cfg文件，同时也记得修改输入的图片大小，保持一致
②对于保存了权重参数和结构的.pt模型
对于这种模型我们可以不用.cfg文件，直接进行转换就行
1）从github里clone一个yolov3.git并解压

git clone -b master https://github.com/ultralytics/yolov3.git

注意：此处的下载的是master分支的文件，注意区分①中下载的archive分支的文件
2）将自己训练完的yolov3-tiny模型copy到解压后的yolov3文件夹中
3）在yolov3的文件夹中运行gen_pt_wts.py即可将.pt模型转换成.wts模型

python3 gen_pt_wts.py -w ./XXX.pt

此处附上gen_pt_wts.py（这个文件不是我写的，是我之前收藏的，目前找不到原链接了，如果原博主看见了，可以直接联系我，我附上引用链接）的内容

import sys
import argparse
import os
import struct
import torch
from utils.torch_utils import select_device


def parse_args():
    parser = argparse.ArgumentParser(description='Convert .pt file to .wts')
    parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
    parser.add_argument('-o', '--output', help='Output (.wts) file path (optional)')
    args = parser.parse_args()
    if not os.path.isfile(args.weights):
        raise SystemExit('Invalid input file')
    if not args.output:
        args.output = os.path.splitext(args.weights)[0] + '.wts'
    elif os.path.isdir(args.output):
        args.output = os.path.join(
            args.output,
            os.path.splitext(os.path.basename(args.weights))[0] + '.wts')
    return args.weights, args.output


pt_file, wts_file = parse_args()

# Initialize
device = select_device('cpu')
# Load model
model = torch.load(pt_file, map_location=device)['model']#.float()  # load to FP32
#print(model)
model.to(device).eval()

with open(wts_file, 'w') as f:
    f.write('{}\n'.format(len(model.state_dict().keys())))
    for k, v in model.state_dict().items():
        vr = v.reshape(-1).cpu().numpy()
        f.write('{} {} '.format(k, len(vr)))
        for vv in vr:
            f.write(' ')
            f.write(struct.pack('>f' ,float(vv)).hex())
        f.write('\n')

至此，模型.pt就可以转换成.wts模型了，可以进行下一步了嘻嘻

【2】.wts模型转.trt/.engine模型

.wts模型转.engine模型比较简单，我在这里按照https://github.com/wang-xinyu/tensorrtx/tree/master/yolov3-tiny里面的步骤详细叙述一下
1）将已经转换好的.wts模型copy到tensorrtx/yolov3-tiny文件夹中
2）执行以下步骤基本上可以完成模型的转换

cd ./tensorrtx/yolov3-tiny
mkdir build
cd build
cmake ..
make
sudo ./yolov3-tiny -s

注意：
（1）自己训练的模型结构要和官方的yolo结构一致，不要在训练的过程中修改模型的结构，不然有些时候转成.wts文件不会报错（譬如采用【1】中的第2）种方式），在转.engine的过程中也会报错，一般会出现缺失某层或者未命名的某层，主要是因为博主的转换程序应该都是按照官方的结构来的，所以一定仔细
（2）在转换的时候记得修改模型的相关参数，主要参数修改的内容和地方如下图所示：

至此，模型的转换就成功了哈哈哈

相关说明
自己训练的yolov3-tiny是指和我一起工作的小伙伴训练的模型，主要与官网训练的yolo模型进行区分说明而已，用自己数据进行训练的过程等我自己实实在在训练过后我再出一篇方法说明

参考链接如下：
【1】https://github.com/wang-xinyu/tensorrtx.git
【2】https://github.com/ultralytics/yolov3.git