开源代码位置在这里，darknet转ONNX模型代码基于python，TensorRT推理代码基于C++。
通过export_onnx.py文件可以将darknet模型转换为ONNX模型，目前可以支持YOLOv3，YOLOv3-SPP，YOLOv4等模型。
parser = argparse.ArgumentParser(description='Transform YOLO weights to ONNX.')
parser.add_argument('--cfg_file', type=str, default='yolov4.cfg', help='yolo cfg file')
parser.add_argument('--weights_file', type=str, default='yolov4.weights', help='yolo weights file')
parser.add_argument('--output_file', type=str, default='yolov4.onnx', help='yolo onnx file')
parser.add_argument('--neck', type=str, default='PAN', help='use which kind neck')
args = parser.parse_args()
cfg_file是darknet的cfg文件的路径，weights_file是darknet的weights文件的路径，output_file是输出的ONNX文件的名称，neck是YOLO的模型结构，YOLOv3相关模型输入FPN，YOLOv4模型相关输入PAN。输入以下命令后，可以将darknet模型转为ONNX模型。
python3 export_onnx.py
导出yolov4-tiny模型的命令是：
python3 export_onnx.py --cfg_file yolov4-tiny.cfg --weights_file yolov4-tiny.weights --output_file yolov4-tiny.onnx --strides 32 16 --neck FPN
转换后的模型加入了transpose，将所有YOLO层的输出结果进行了concat，最终结构如下：
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等。
Yolov4:
onnx_file: "../yolov4.onnx"
engine_file: "../yolov4.trt"
labels_file: "../coco.names"
BATCH_SIZE: 1
INPUT_CHANNEL: 3
IMAGE_WIDTH: 608
IMAGE_HEIGHT: 608
obj_threshold: 0.4
nms_threshold: 0.45
stride: [8, 16, 32]
anchors: [[12, 16], [19, 36], [40, 28], [36, 75], [76, 55], [72, 146], [142, 110], [192, 243], [459, 401]]
通过以下命令对项目进行编译，生成Yolov4_trt
mkdir build && cd build
cmake ..
make -j
通过以下命令运行项目，得到推理结果
./Yolov4_trt ../config.yaml ../samples
推理结果如下图所示：
可以看到输出结果与darknet模型的输出结果包括置信度完全一致，在输出的信息中，每张图片的平均处理时间约为16ms，单独engine的推理时间约为7ms。