使用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，本文的核心思想包括：
上面是250fps的人脸检测模型，得益于TensorRT的加速。输入尺寸为1280x960.
为什么要用onnx
现在大家都喜欢用pytorch训练模型，而pytorch训练的模型转成pth，用C++推理也很难达到真正的加速效果，因为本质上最耗时的网络前向推理部分并没有太多的加速。并且采用libtorch C++推理pytorch并不是一件简单的事情，除非你的模型可以被trace。
在这种情况之下，引入onnx更合理，从目前整个DL生态来看，onnx具有以下好处：
前段时间，我们release了一个retinaface的pytorch项目，并且我们想办法将它导出到了onnx模型，当然这期间经过一些修改，没有复杂模型的代码可以在不修改的情况下轻而易举export到onnx，关于这部分代码可以在我们的平台上找到:
我们今天要做的事情，就是在上面的onnx模型的基础上，采用TensorRT来进行推理。先做一个简单的速度对比：
耗时(s)可以看到，采用TensorRT对ONNX模型加速，速度提升可以说是天囊之别。并且，采用TensorRT纯C++推理可以在语言层面获得更多的加速。 我们实现的TensorRT加速的Retinaface应该是目前来讲面向GPU速度最快的检测方案，并且可以同时生成bbox和landmark，相比于MTCNN，模型更加简单，推理更加快速，准确度更高 .
真正落地的算法部署，毫无疑问，假如你的target是GPU，采用ONNX+TensorRT应该是目前最成熟、最优化的方案。假如你的target是一些嵌入式芯片，那么采用MNN也是可以通过onnx轻松实现CPU嵌入式端快速推理的。
既然ONNX和TensorRT这么好，为什么都不用呢？为什么都还在用Python写难看的推理的代码呢？原因也很简单：
今天这篇教程便是教大家如何一步一步的实现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。
在这里，假如你是新手，有一点需要注意：
闲话不多说，假如我们拿到了trt的engine，我们如何进行推理呢？总的来说，分为3步：
ICudaEngine
, 这个是TensorRT推理的核心；
当然这里最核心的东西其实就两个，一个是如何导入拿到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平台上转到并下载：
大家可以看到，假如你想对模型进行进一步的加速，实际上也是在这上面进行。当你拿到你的
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的训练代码可以在这里找到：
TensorRT部署完整的代码可以在这里找到：
我们看到，随着AI技术的不断成熟，大家已经不局限于用简单的python来编写古老的代码，我们致力于寻找更前沿的AI部署方案，TensorRT就是其中的一种，我们发现，通过对网络模型本身的思考优化、通过对我们网络计算框架的思考和优化、通过对编写网络推理语言和算法本身的思考和优化，构建了一道深不可测的技术瓶颈和壁垒。未来大家可能会看到，为什么你的MaskRCNN只有10fps，而别人的可以在全尺寸（1280p）下跑到35fps？
方寸之间，尽显功夫。
未来我们将继续在onnx-tensorrt的技术路线为大家奉献更加高质量的代码，我们的下一个目标是采用ONNX推理，并且用TensorRT加速MaskRCNN。Detectron2都发布了，这个还会远吗？Instance segmentation 和全景分割的Realtime inference是我们的终极目标！
其实看完这篇文章，建议大家可以做的事情：
当然，欢迎大家评论和转发，我们有机会也会开源我们踩坑之后的收获。
我们为维护、编写、创造这些代码花费了许多宝贵的时间，同时维护他们也需要海量的云平台，我们致力于帮助更多的初学者、中级学者、老司机提供完善的AI代码部署平台，如果你对AI感兴趣，可以通过我们的论坛来交流。
另外我们开通了Slack群，欢迎大家来交流：
