查看原文
其他

使用ONNX+TensorRT部署人脸检测和关键点250fps

金天 极市平台 2021-09-20

加入极市专业CV交流群,与6000+来自腾讯,华为,百度,北大,清华,中科院等名企名校视觉开发者互动交流!更有机会与李开复老师等大牛群内互动!

同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总,行业技术交流。关注 极市平台 公众号 ,回复 加群,立刻申请入群~


作者:金天

来源:https://zhuanlan.zhihu.com/p/86827710

本文已经作者授权,禁止二次转载


我尽量用尽可能短的语言将本文的核心内容浓缩到文章的标题中,前段时间给大家讲解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,关于这部分代码可以在这个平台上找到 : http://manaai.cn/


我们今天要做的事情,就是在上面的onnx模型的基础上,采用TensorRT来进行推理。先做一个简单的速度对比:

可以看到,采用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模型可以从这个地方下载:

AI算法详情页 ( 链接: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



未来计划


我们看到,随着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重新训练一个手和关键点的检测,实现手的姿态检测。



-End-



*延伸阅读



人脸检测交流群


添加极市小助手微信(ID : cv-mart),备注:研究方向-姓名-学校/公司-城市(如:目标检测-小极-北大-深圳),即可进群。(已经添加小助手的好友直接私信)


△长按添加极市小助手


其他方向如目标检测、目标跟踪、工业检测、医学影像、三维&SLAM、图像分割等,也可扫码添加极市小助手拉你进群,更有每月大咖直播分享、真实项目需求对接、干货资讯汇总,行业技术交流,一起来让思想之光照的更远吧~

△长按关注极市平台


觉得有用麻烦给个在看啦~  

: . Video Mini Program Like ,轻点两下取消赞 Wow ,轻点两下取消在看

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存