CenterNet101 TensorRT加速

本文介绍 CenterNet101 TensorRT加速

CenterNet101 TensorRT加速

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

最近我们发现CenterNet在做目标检测的时候具有十分优良的特性,比如它可以很好的和多任务进行组合,速度快,同时给出的检测结果在概率上更为可信等等。这是在实际对比不同检测算法上得到的结论,如果你考虑在实际产品或者工业应用上使用检测算法,那么在足够的算力下,最好不要采用一阶段算法,无数案例表明,一阶段算法在产生误检的概率上更大。

同时CenterNet也带来一个问题,那就是部署不太容易,主要是两个方面:

  • 主流实现大多不好支持onnx导出;
  • 后处理与传统的检测算法不太一样,比如nms,CenterNet用的实际上是一个3x3的maxpooling。

我们早在很久以前就将CenterNet导出为了onnx模型,现在我们已经实现了使用TensorRT对CenterNet进行加速。先看看加速的结果:

大概在GTX1070上可以达到35ms,GTX1080TI上达到22ms,最快可以加速到50FPS。采用int8推理速度可以更快。其实在一个R101的backbone下,速度能达到22ms,50fps已经很快了。相比于YoloV3+Darknet53SPP,其TensorRT加速后速度大概为40ms左右,显然CenterNet101能做到速度更快,精度更高(均为GTX1080ti下)。

本篇文章不仅仅只是告诉大家,CenterNet可以用TensorRT加速,同时也想告诉大家整个过程大概是怎么样的,在文章末尾会给出整个工程代码链接。另外,本文需要支持导出onnx的CenterNet版本,如果大家想直接下载的话,可以从我们的平台下载:

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

## CenterNet101加速过程

其实在很早之前,我就在MANA社区提出了一个关于CenterNet后处理的思考,链接. 主要是onnx到TensorRT的两个问题:

  • 后处理本质上是支持把这些操作同时trace到onnx模型的,它甚至不需要单独的写Plugin,因为它只需要两个onnx的op,即maxpool和topk;
  • 但问题是这两个op,在TensorRT中并没有相应的支持。

最终我们选择onnx的导出方式是,只导出网络部分,对于后处理部分,我们再单独的写layer来支持它。不过在最开始这一步是很难完成的,原因主要是:

  • 后处理有点复杂,虽然从原理上来讲,只是一个对 1x80x128x128的特征热力图的maxpool,但是自己写代码就有点繁杂了,最关键的是速度可能会很慢;
  • 直接从网络拿到输出,再自己写代码后处理会涉及到一个topK的问题,即找到张量里面特定维度的最大值的index,这个要自己手鲁代码估计会非常非常的慢。

好在最近有社区朋友放出一个CUDA Kernel很好的解决了这个问题,可以说这个kernel就是本文实现CenterNet TensorRT加速的核心。这个kernel的核心逻辑,其实就是解决了我上面提到的后处理的问题。

这部分CUDA代码贴出来其实也十分的简单:

__global__ void Cdetforward_kernel(const float *hm, const float *reg,const float *wh ,
float *output,const int w,const int h,const int classes,const int kernerl_size,const float visthresh ) {
int idx = (blockIdx.x + blockIdx.y * gridDim.x) * blockDim.x + threadIdx.x;
if (idx >= w*h) return;
int padding = kernerl_size/2;
int offset = - padding /2;
int stride = w*h;
int grid_x = idx % w ;
int grid_y = idx / w ;
int cls,l,m;
float c_x,c_y;
for (cls = 0; cls < classes; ++cls )
{
int objIndex = stride * cls + idx;
float objProb = hm[objIndex];
float max=-1;
int max_index =0;
for(l=0 ;l < kernerl_size ; ++l)
for(m=0 ; m < kernerl_size ; ++m){
int cur_x = offset + l + grid_x;
int cur_y = offset + m + grid_y;
int cur_index = cur_y * w + cur_x + stride*cls;
int valid = (cur_x>=0 && cur_x < w && cur_y >=0 && cur_y <h );
float val = (valid !=0 ) ? Logist(hm[cur_index]): -1;
max_index = (val > max) ? cur_index : max_index;
max = (val > max ) ? val: max ;
}
objProb = Logist(objProb);
if((max_index == objIndex) && (objProb > visthresh)){
int resCount = (int)atomicAdd(output,1);
//printf("%d",resCount);
char* data = (char * )output + sizeof(float) + resCount*sizeof(Detection);
Detection* det = (Detection*)(data);
c_x = grid_x + reg[idx] ; c_y = grid_y + reg[idx+stride];
det->bbox.x1 = (c_x - wh[idx]/2)*4;
det->bbox.y1 = (c_y - wh[idx+stride]/2)*4 ;
det->bbox.x2 = (c_x + wh[idx]/2)*4;
det->bbox.y2 = (c_y + wh[idx+stride]/2)*4;
det->classId = cls;
det->prob = objProb;
}
}
}

需要注意的是,这里面的kernelsize就是maxpool的边长,通过设定不同的值你可以控制整个网络的粒度。最后我们通过一个threshold在这部分里面踢掉了一些置信度比较低的heappoint。

最后将结果保存为一个Detection 的数据结构中。

就这么简单,这边是CenterNet TensorRT加速的核心。最终我们也实现了通过这个后处理可以完美的得到和python一样的结果,同时拥有更快的推理速度。

应用

CenterNet在应用前景不仅仅是目标检测上,它的网络结构比传统的基于anchor的检测方法更加简单,比基于RPN的二阶段方法更是简单的过于优雅,更重要的是它还可以用来组合不同的任务,比如centerface就是将CenterNet和face landmark组合生成的网络模型,在单目的3D模型中,CenterNet也能够很好的预测基于单目的3D回归。

Future Work

未来将进一步添加更轻量级后端的CenterNet支持和对应的TensorRT加速,使得模型更加适用于边缘计算。最终将算法落地部署打通。

对于使用CenterNet训练自己的数据,并想用TensorRT进行加速部署的同学,用这一套就已经足够了,剩下的工作,只需要改变你的类别数即可了。

本文中设计到的代码包括:

CenterNet TensorRT加速

支持ONNX模型导出的CenterNet版本,pytorch1.3支持

最后,如果你对人工智能感兴趣,或者你正在致力于做AI算法的落地,欢迎加入我们的社区来进行交流:

t.manaai.cn

往期模型推荐

再次推荐我们上一期的模型:TensorRT加速的Retinaface,速度高达250fps!

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