最快的人臉檢測！ONNX+TensorRT只要4ms!

使用ONNX+TensorRT部署人臉檢測和關鍵點250fps

This article was original written by Jin Tian, welcome re-post, first come with jinfagang.github.io . but please keep this copyright info, thanks, any question could be asked via wechat: jintianiloveu

我儘可能用盡量短的語言將本文的核心內容濃縮到文章的標題中，前段時間給你們講解Jetson Nano的部署，咱們講到用caffe在Nano上部署yolov3，感興趣的童鞋能夠看看以前的文章，而後順便挖了一個坑：如何部署ONNX模型, 這個問題其實分爲兩個部分，第一是爲何要用ONNX，第二是如何部署ONNX。本文就是來填這個坑的。

TLTR，本文的核心思想包括：

• 怎麼才能最快速的部署ONNX模型；
• 當今世界檢測人臉和關鍵點最快的模型是什麼？
• 如何使用ONNX+TensorRT來讓你的模型提高7倍加速
• 咱們將向你們介紹咱們的新一代人臉檢測+比對識別的新一代引擎，有望在GPU上跑到200fps以上，固然也將開源。
• 如何使用C++在TensorRT上部署ONNX模型。

上面是250fps的人臉檢測模型，得益於TensorRT的加速。輸入尺寸爲1280x960.

爲何要用onnx

如今你們都喜歡用pytorch訓練模型，而pytorch訓練的模型轉成pth，用C++推理也很難達到真正的加速效果，由於本質上最耗時的網絡前向推理部分並無太多的加速。而且採用libtorch C++推理pytorch並非一件簡單的事情，除非你的模型能夠被trace。

在這種狀況之下，引入onnx更合理，從目前整個DL生態來看，onnx具備如下好處：

• 它的模型格式比基於layer的老一代框架更加細粒度；
• 它擁有統一化的定義，而且基於任何框架均可以推理他；
• 它能夠實現不一樣框架之間的互相轉化。

前段時間，咱們release了一個retinaface的pytorch項目，而且咱們想辦法將它導出到了onnx模型，固然這期間通過一些修改，沒有複雜模型的代碼能夠在不修改的狀況下垂手可得export到onnx，關於這部分代碼能夠在咱們的平臺上找到:

manaai.cn

咱們今天要作的事情，就是在上面的onnx模型的基礎上，採用TensorRT來進行推理。先作一個簡單的速度對比：

框架	語言	耗時(s)	fps
Pytorch	python	0.012+0.022	29
ONNXRuntime	python	0.008+0.022	34
TensorRT	C++	0.004+0.001	250

能夠看到，採用TensorRT對ONNX模型加速，速度提高能夠說是天囊之別。而且，採用TensorRT純C++推理能夠在語言層面得到更多的加速。咱們實現的TensorRT加速的Retinaface應該是目前來說面向GPU速度最快的檢測方案，而且能夠同時生成bbox和landmark，相比於MTCNN，模型更加簡單，推理更加快速，準確度更高.

真正落地的算法部署，毫無疑問，假如你的target是GPU，採用ONNX+TensorRT應該是目前最成熟、最優化的方案。假如你的target是一些嵌入式芯片，那麼採用MNN也是能夠經過onnx輕鬆實現CPU嵌入式端快速推理的。

既然ONNX和TensorRT這麼好，爲何都不用呢？爲何都還在用Python寫難看的推理的代碼呢？緣由也很簡單：

• 入門太難，C++通常人玩不來；
• 既要懂模型的每一層輸入輸出，又要懂TensorRT的API，至少要很熟悉。

今天這篇教程即是教你們如何一步一步的實現TensorRT實現最快速的推理。先來看看實際TensorRT加速的效果：

看圖片看不出啥來，看視頻：

效果仍是很是不錯的。

Retinaface模型簡單介紹

retinaface是Insightface作的一個動做(DeepInsight), 可是原始的只有MXNet版本，這個網絡模型具備小巧精度高特色，而且它是一個帶有landmark分支輸出的網絡，這使得該模型能夠輸出landmark。

這個網絡之因此叫作retina是由於它引入了FPN的結構和思想，使得模型在小尺度的臉上具備更好的魯棒性。

在這裏咱們引入一個工具：sudo pip3 install onnxexplorer 能夠快速的查看咱們的onnx模型的結構，咱們須要用到的onnx模型能夠從這個地方下載：manaai.cn/aicodes_det…

咱們作了一些修改使得pytorch的模型能夠導出到onnx，而且咱們作了一些特殊的處理，使得onnx模型能夠經過 onnx2trt轉到TensorRT的engine。

TensorRT C++推理

接下來應該是本文的核心內容了，上面提到的 onnx2trt 能夠經過編譯 https://gitub.com/onnx/onnx-tensorrt 倉庫，來獲得 onnx2trt ，經過這個執行程序，能夠將onnx轉到trt的engine。

在這裏，假如你是新手，有一點須要注意：

• 並非全部的onnx都可以成功轉到trt engine，除非你onnx模型裏面全部的op都被支持；
• 你須要在電腦中安裝TensorRT 6.0，由於只有TensorRT6.0支持動態的輸入。

閒話很少說，假如咱們拿到了trt的engine，咱們如何進行推理呢？總的來講，分爲3步：

1. 首先load你的engine，拿到一個 ICudaEngine, 這個是TensorRT推理的核心；
2. 你須要定位你的模型的輸入和輸出，有幾個輸入有幾個輸出；
3. forward模型，而後拿到輸出，對輸出進行後處理。

固然這裏最核心的東西其實就兩個，一個是如何導入拿到CudaEngine，第二個是比較麻煩的後處理。

IBuilder* builder = createInferBuilder(gLogger);
  assert(builder != nullptr);
  nvinfer1::INetworkDefinition* network = builder->createNetwork();
  auto parser = nvonnxparser::createParser(*network, gLogger);

  if ( !parser->parseFromFile(modelFile.c_str(), static_cast<int>(gLogger.reportableSeverity) ) )
  {
    cerr << "Failure while parsing ONNX file" << std::endl;
  }


  IHostMemory *trtModelStream{nullptr};
  // Build the engine
  builder->setMaxBatchSize(maxBatchSize);
  builder->setMaxWorkspaceSize(1 << 30);

  if (mTrtRunMode == RUN_MODE::INT8) {
    std::cout << "setInt8Mode" << std::endl;
    if (!builder->platformHasFastInt8())
      std::cout << "Notice: the platform do not has fast for int8" << std::endl;
// builder->setInt8Mode(true);
// builder->setInt8Calibrator(calibrator);
    cerr << "int8 mode not supported for now.\n";
  } else if (mTrtRunMode == RUN_MODE::FLOAT16) {
    std::cout << "setFp16Mode" << std::endl;
    if (!builder->platformHasFastFp16())
      std::cout << "Notice: the platform do not has fast for fp16" << std::endl;
    builder->setFp16Mode(true);
  }

  ICudaEngine* engine = builder->buildCudaEngine(*network);
  assert(engine);
  // we can destroy the parser
  parser->destroy();
  // serialize the engine, then close everything down
  trtModelStream = engine->serialize();
  trtModelStream->destroy();
  InitEngine();
複製代碼

這個是咱們維護的 onnx_trt_engine的一部分，這段代碼的做用是直接將你以前生成的trt engine，導入到你的ICudaEngine之中。你們若是須要完整的code，能夠在咱們的MANA平臺上轉到並下載：

manaai.cn

你們能夠看到，假如你想對模型進行進一步的加速，實際上也是在這上面進行。當你拿到你的 iCudaEngine以後，剩下的事情就是根據你的model的output name拿到對應的輸出。整個過程其實仍是能夠一鼓作氣的，惟一可能複雜一點的是你須要動態allocate對應大小size的data。

auto out1 = new float[bufferSize[1] / sizeof(float)];
  auto out2 = new float[bufferSize[2] / sizeof(float)];
  auto out3 = new float[bufferSize[3] / sizeof(float)];

  cudaStream_t stream;
  CHECK(cudaStreamCreate(&stream));
  CHECK(cudaMemcpyAsync(buffers[0], input, bufferSize[0], cudaMemcpyHostToDevice, stream));
// context.enqueue(batchSize, buffers, stream,nullptr);
  context.enqueue(1, buffers, stream, nullptr);

  CHECK(cudaMemcpyAsync(out1, buffers[1], bufferSize[1], cudaMemcpyDeviceToHost, stream));
  CHECK(cudaMemcpyAsync(out2, buffers[2], bufferSize[2], cudaMemcpyDeviceToHost, stream));
  CHECK(cudaMemcpyAsync(out3, buffers[3], bufferSize[3], cudaMemcpyDeviceToHost, stream));
  cudaStreamSynchronize(stream);

  // release the stream and the buffers
  cudaStreamDestroy(stream);
  CHECK(cudaFree(buffers[0]));
  CHECK(cudaFree(buffers[1]));
  CHECK(cudaFree(buffers[2]));
  CHECK(cudaFree(buffers[3]));
複製代碼

這是如何從TensorRT推理的結果轉到咱們的CPU上來，而且經過Async來同步數據，最終你拿到的數據將在你事先定義好的buffer裏面，再進行後處理便可。

因爲C++代碼過於龐大和複雜，這些代碼將會開源到咱們的MANA AI平臺。固然咱們花費了不少力氣來編寫教程，而且提供源碼，若是你對AI感興趣，而缺少一個好的學習羣體和導師，不妨加入咱們的會員計劃，咱們是一個致力於打造工業級前沿黑科技的AI學習者羣體。

咱們pytorch的訓練代碼能夠在這裏找到：

manaai.cn/aicodes_det…

TensorRT部署完整的代碼能夠在這裏找到：

manaai.cn/aicodes_det…

將來計劃

咱們看到，隨着AI技術的不斷成熟，你們已經不侷限於用簡單的python來編寫古老的代碼，咱們致力於尋找更前沿的AI部署方案，TensorRT就是其中的一種，咱們發現，經過對網絡模型自己的思考優化、經過對咱們網絡計算框架的思考和優化、經過對編寫網絡推理語言和算法自己的思考和優化，構建了一道深不可測的技術瓶頸和壁壘。將來你們可能會看到，爲何你的MaskRCNN只有10fps，而別人的能夠在全尺寸（1280p）下跑到35fps？

方寸之間，盡顯功夫。

將來咱們將繼續在onnx-tensorrt的技術路線爲你們奉獻更加高質量的代碼，咱們的下一個目標是採用ONNX推理，而且用TensorRT加速MaskRCNN。Detectron2都發布了，這個還會遠嗎？Instance segmentation 和全景分割的Realtime inference是咱們的終極目標！

其實看完這篇文章，建議你們能夠作的事情：

1. 把咱們的TensorRT版本retinaface部署到Jetson Nano上，你能夠獲得一個至少在30fps的人臉檢測模型；
2. 你能夠嘗試用Retinaface從新訓練一個手和關鍵點的檢測，實現手的姿態檢測。

固然，歡迎你們評論和轉發，咱們有機會也會開源咱們踩坑以後的收穫。

咱們爲維護、編寫、創造這些代碼花費了許多寶貴的時間，同時維護他們也須要海量的雲平臺，咱們致力於幫助更多的初學者、中級學者、老司機提供完善的AI代碼部署平臺，若是你對AI感興趣，能夠經過咱們的論壇來交流。

另外咱們開通了Slack羣，歡迎你們來交流：

join.slack.com/t/manaai/sh…