如何將自己訓練的yolo模型轉為tensorrt相關的模型

• 【1】.pt模型轉.wts模型
• 【2】.wts模型轉.trt/.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 formatmodel.load_state_dict(torch.load(weights, map_location=device)['model']) else: # darknet formatload_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_devicedef 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.outputpt_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

總結

以上是生活随笔為你收集整理的2022-10-31如何将自己训练的yolo模型转为tensorrt相关的模型（.trt or .engine）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。