CenterFace模型转TensorRT

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\times 32$的，因为TensorRT推理引擎的输入都是固定的，所以需要修改ONNX模型的输入和输出尺寸。如下图所示，输入和输出的分辨率修改为$640\times 640$和$160\times 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. 推理结果

推理结果如下图所示：