将darknet格式的权重转为onnx格式,然后再转为tensorrt格式
首先准备好一个darknet格式的权重文件(.weights),配置文件(.cfg),训练集名称文件(.names)
一、darknet转为onnx
这里我使用的是这个开源链接
https://github.com/Tianxiaomo/pytorch-YOLOv4
通过一个python脚本,将darknet格式的权重,转为onnx格式,并且可以输入一张图片,这个脚本会对这张图片进行识别 具体的步骤为: (1)先把上面这个github里面的文件下载下来 (2)运行命令:
python demo_darknet2onnx.py yolov4-tiny-416.cfg obj.names yolov4-tiny-416.weights fire1.jpg 1
按照上面图1的格式来写这个命令就行
运行完成之后,就会生成一个.onnx格式的权重了。
二、将onnx格式转为tensorrt格式
这里也是参考上面github里面的
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
tensorrt格式的权重好像有两种,一种是.trt,一种是.engine。我这里转化的是.trt格式,.engine格式没试过。
运行命令:
/opt/TensorRT-8_0_1_6/bin/trtexec --onnx=yolov4_1_3_416_416_static.onnx --explicitBatch --saveEngine=yolov4-tiny-416.trt --workspace=2048 --fp16
(这个转格式时间有点长,可能得个10分钟)
备注:
如果提示没有trtexec,需要安装tensorrt,然后在tensorrt的安装文件夹下面的samples里面
cd <TensorRT root directory>/samples/trtexec
make
编译之后,在<TensorRT root directory>/bin
里面就有trtexec了。
参考链接:
https://github.com/Tianxiaomo/pytorch-YOLOv4