jetson系列开发板生成.engine部署

原创已于 2023-07-18 13:05:45 修改 · 1.2k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#python #pytorch #视觉检测 #边缘计算

于 2023-07-17 11:03:51 首次发布

Deep Learning 同时被 2 个专栏收录

14 篇文章

订阅专栏

日常问题记录

10 篇文章

订阅专栏

该文介绍了如何在Jetson系列开发板上使用TensorRTx的Yolov5实现模型转换，包括从训练好的.pt文件生成.wts权重文件，然后在TensorRTx的构建目录下编译并生成Engine，用于加速推理过程。步骤涉及下载TensorRTx源码、修改模型参数、编译项目以及使用Python脚本转换权重等操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

jetson系列开发板生成engine

1. 下载 tensorrtx 至 NX

这个tensorrtx的版本号要和训练时的版本号一样。

git clone -b yolov5-v5.0 https://github.com/wang-xinyu/tensorrtx.git

2. 修改 yololayer.h 中的参数

$[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-76KPGlS4-1689562839272)(C:\Users\Administrator\AppData\Roaming\Typora\typora-user-images\image-20221025185746573.png)]$

训练时有几类就将 CLASS_NUM 修改为几。

static constexpr int CLASS_NUM = 80;

3. 编译 tensorrtx/yolov5

cd tensorrtx/yolov5
mkdir build
cd build
cmake ..
make

4. 在训练yolov5的平台上生成 .wts文件

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)  # load to FP32
model = model['ema' if model.get('ema') else 'model'].float()

# update anchor_grid info
anchor_grid = model.model[-1].anchors * model.model[-1].stride[...,None,None]
# model.model[-1].anchor_grid = anchor_grid
delattr(model.model[-1], 'anchor_grid')  # model.model[-1] is detect layer
model.model[-1].register_buffer("anchor_grid",anchor_grid) #The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight.

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')

命令行执行

python gen_wts.py -w xxxx.pt -o xxxx.wts

4. copy .wts 文件到tensorrtx/yolov5/build目录下

5. 在build目录下执行下面语句生成 engine

sudo ./yolov5 -s hat.wts hat.engine s //需要注意的是最后一个参数s，这个参数表示是用yolov5的哪个模型训练的（s,m,l...)