文章目录
yolov7 PyTorch模型转TensorRT
1. github开源代码
yolov7 TensorRT推理的开源代码位置在https://github.com/linghu8812/tensorrt_inference/tree/master/project/yolov7，yolov7官方的开源代码位置为https://github.com/WongKinYiu/yolov7，对yolov7 epxort ONNX模型的代码作了少许修改，主要是(1) 在导出ONNX模型时，增加了onnxsim进行模型简化；(2) 在导出ONNX模型时进行了少许修改，可以去掉ONNX模型的ScatterND的op，使得TensorRT可以直接解析ONNX模型。
2. PyTorch模型转ONNX模型
首先通过命令 git clone https://github.com/linghu8812/yolov7.gitclone yolov7的代码，通过以下命令生成ONNX文件。 --weights可以指定模型文件路径， --simplify进行ONNX模型简化， --grid使输出为一个tensor， --img-size为输入图片尺寸， --batch-size设置batch size的大小。
python export.py --weights ./weights/yolov7.pt --simplify --grid
对于1280输入的模型，可通过以下命令生成ONNX模型：
python export.py --weights ./weights/yolov7-w6.pt --simplify --grid --img-size 1280
3. ONNX模型转TensorRT模型
3.1 概述
TensorRT模型即TensorRT的推理引擎，代码中通过C++实现。相关配置写在config.yaml文件中，如果存在 engine_file的路径，则读取 engine_file，否则从 onnx_file生成 engine_file。
void yolov7::LoadEngine() {
std::fstream existEngine;
existEngine.open(engine_file, std::ios::in);
if (existEngine) {
readTrtFile(engine_file, engine);
assert(engine != nullptr);
} else {
onnxToTRTModel(onnx_file, engine_file, engine, BATCH_SIZE);
assert(engine != nullptr);
}
}
config.yaml文件可以设置batch size，图像的size，关于anchor的处理已经写到了ONNX模型中，所以处理时只需知道每个检测层的stride和anchor的数量即可，不需要再配置具体的anchor。
yolov7:
onnx_file: "../yolov7.onnx"
engine_file: "../yolov7.trt"
labels_file: "../coco.names"
BATCH_SIZE: 1
INPUT_CHANNEL: 3
IMAGE_WIDTH: 640
IMAGE_HEIGHT: 640
obj_threshold: 0.4
nms_threshold: 0.45
strides: [8, 16, 32]
num_anchors: [3, 3, 3]
yolov7-w6模型的配置文件如下：
yolov7:
onnx_file: "../yolov7-w6.onnx"
engine_file: "../yolov7-w6.trt"
labels_file: "../coco.names"
BATCH_SIZE: 1
INPUT_CHANNEL: 3
IMAGE_WIDTH: 1280
IMAGE_HEIGHT: 1280
obj_threshold: 0.4
nms_threshold: 0.45
strides: [8, 16, 32, 64]
num_anchors: [3, 3, 3, 3]
3.2 编译
通过以下命令对项目进行编译，生成 yolov7_trt
mkdir build && cd build
cmake ..
make -j
3.3 运行
通过以下命令运行项目，得到推理结果
./yolov7_trt ../config.yaml ../samples
4. 推理结果
推理结果如下图所示：
Original: https://blog.csdn.net/linghu8812/article/details/125741951
Author: 令狐傻笑
Title: yolov7 PyTorch模型转TensorRT
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/651544/
转载文章受原作者版权保护。转载请注明原作者出处！