最快的人脸检测！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…