文章目录
yolov5 PyTorch模型转TensorRT
1. github开源代码
yolov5 TensorRT推理的开源代码位置在https://github.com/linghu8812/tensorrt_inference/tree/master/yolov5,PyTorch转onnx的代码见从原作者fork过来的代码:https://github.com/linghu8812/yolov5,对模型转换做了一点修改。
2. PyTorch模型转ONNX模型
首先通过命令git clone https://github.com/linghu8812/yolov5.git
clone yolov5的代码,然后将export_onnx.py文件拷贝到yolov5/models
文件夹中,通过以下命令生成ONNX文件。对于yolov5s,yolov5m,yolov5l,yolov5x这几个模型都可以支持。--weights
可以指定模型文件路径,--img
为输入图片尺寸,--batch
设置batch size的大小。
export PYTHONPATH="$PWD" && python3 export_onnx.py --weights ./weights/yolov5s.pt --img 640 --batch 1
同YOLOv4模型一样,对输出结果也做了concat,如下图所示。
代码中加入了onnxsim模块,以简化后续的tranpose和reshape节点的结构,简化代码如下:
onnx_model = onnx.load(f) # load onnx model
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, f)
模型结构简化为:
3. ONNX模型转TensorRT模型
3.1 概述
TensorRT模型即TensorRT的推理引擎,代码中通过C++实现。相关配置写在config.yaml文件中,如果存在engine_file
的路径,则读取engine_file
,否则从onnx_file
生成engine_file
。
void YOLOv5::LoadEngine() {
// create and load engine
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等。
yolov5:
onnx_file: "../yolov5x.onnx"
engine_file: "../yolov5x.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
stride: [8, 16, 32]
anchors: [[10,13], [16,30], [33,23], [30,61], [62,45], [59,119], [116,90], [156,198], [373,326]]
3.2 编译
通过以下命令对项目进行编译,生成yolov5_trt
mkdir build && cd build
cmake ..
make -j
3.3 运行
通过以下命令运行项目,得到推理结果
./yolov5_trt ../config.yaml ../samples
4. 推理结果
推理结果如下图所示:
上图为yolov5x的测试结果,在输出的信息中,yolov5x每张图片的平均处理时间约为22.9ms,单独engine的推理时间约为12.1ms。