文章目录
CenterFace模型转TensorRT
1. github开源代码
CenterFaceTensorRT推理的开源代码位置在https://github.com/linghu8812/tensorrt_inference/tree/master/CenterFace,作者开源的代码位置在https://github.com/Star-Clouds/CenterFace,论文的arxiv地址为https://arxiv.org/abs/1911.03599。
2. 重写ONNX模型
作者在github上放了两个开源模型,在转换TensorRT时,通过日志可以看到这两个模型是通过PyTorch转换出来的,作者没有开源通过PyTorch构建模型的代码,也没有开源PyTorch版的模型。通过netron可以看到作者开源的ONNX模型的分辨率为 32 × 32 32 \times 32 32×32的,因为TensorRT推理引擎的输入都是固定的,所以需要修改ONNX模型的输入和输出尺寸。如下图所示,输入和输出的分辨率修改为 640 × 640 640\times 640 640×640和 160 × 160 160\times 160 160×160
python3 export_onnx.py
3. ONNX模型转TensorRT模型
3.1 概述
TensorRT模型即TensorRT的推理引擎,代码中通过C++实现。相关配置写在config.yaml文件中,如果存在engine_file
的路径,则读取engine_file
,否则从onnx_file
生成engine_file
。
void CenterFace::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和检测的threshold即可,这是anchor free模型相比于需要anchor进行回归的模型的优势。
CenterFace:
onnx_file: "../centerface.onnx"
engine_file: "../centerface.trt"
BATCH_SIZE: 1
INPUT_CHANNEL: 3
IMAGE_WIDTH: 640
IMAGE_HEIGHT: 640
obj_threshold: 0.5
nms_threshold: 0.45
在将图像的原始数据转换为张量时,需要保持图像的长宽比,相应的代码如下
float ratio = float(IMAGE_WIDTH) / float(src_img.cols) < float(IMAGE_HEIGHT) / float(src_img.rows) ? float(IMAGE_WIDTH) / float(src_img.cols) : float(IMAGE_HEIGHT) / float(src_img.rows);
cv::Mat flt_img = cv::Mat::zeros(cv::Size(IMAGE_WIDTH, IMAGE_HEIGHT), CV_8UC3);
cv::Mat rsz_img;
cv::resize(src_img, rsz_img, cv::Size(), ratio, ratio);
rsz_img.copyTo(flt_img(cv::Rect(0, 0, rsz_img.cols, rsz_img.rows)));
flt_img.convertTo(flt_img, CV_32FC3);
3.2 编译
通过以下命令对项目进行编译,生成yolov5_trt
mkdir build && cd build
cmake ..
make -j
3.3 运行
通过以下命令运行项目,得到推理结果
./CenterFace_trt ../config.yaml ../samples
4. 推理结果
推理结果如下图所示: