ONNX+TensorRT

 

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

This article was original written by Jin Tian, welcome re-post, first come with https://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模型。

image

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

爲什麼要用onnx

現在大家都喜歡用pytorch訓練模型，而pytorch訓練的模型轉成pth，用C++推理也很難達到真正的加速效果，因爲本質上最耗時的網絡前向推理部分並沒有太多的加速。並且採用libtorch C++推理pytorch並不是一件簡單的事情，除非你的模型可以被trace。

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

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

前段時間，我們release了一個retinaface的pytorch項目，並且我們想辦法將它導出到了onnx模型，當然這期間經過一些修改，沒有複雜模型的代碼可以在不修改的情況下輕而易舉export到onnx，關於這部分代碼可以在我們的平臺上找到:

http://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加速的效果：

image

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

image

效果還是非常不錯的。

Retinaface模型簡單介紹

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

這個網絡之所以叫做retina是因爲它引入了FPN的結構和思想，使得模型在小尺度的臉上具有更好的魯棒性。

在這裏我們引入一個工具：sudo pip3 install onnxexplorer 可以快速的查看我們的onnx模型的結構，我們需要用到的onnx模型可以從這個地方下載：http://manaai.cn/aicodes_detail3.html?id=46

我們做了一些修改使得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平臺上轉到並下載：

http://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的訓練代碼可以在這裏找到：

http://manaai.cn/aicodes_detail3.html?id=46

TensorRT部署完整的代碼可以在這裏找到：

http://manaai.cn/aicodes_detail3.html?id=48