TensorRT+CUDA加速优化版CenterNet旋转目标以及水平目标框的检测
前言
由于工作项目所需，一直用centerNet做旋转目标检测，在实际产品或者工业应用上落地此检测算法，那么在足够的算力下， 更好优选的方式还是需要c/c++来部署实现。
那么CenterNet也带来一个问题，那就是部署不太容易，主要是两个方面：
[*]主流实现大多不好支持onnx导出；
[*]后处理与传统的检测算法不太一样，比如nms，CenterNet用的实际上是一个3x3的maxpooling。
[*]此处还涉及到一点就是，原版的centerNet并可以检测旋转目标，所以此处就涉及到一个角度问题。
首选是将centerNet导出为onnx模型，之后再使用TensorRT对CenterNet进行加速。
本篇文章不仅仅只是告诉大家，CenterNet可以用TensorRT+cuda加速，同时也想告诉大家整个过程大概是怎么样的，最终我也会将我的整个工程代码链接放出来。（本文需要支持导出onnx的CenterNet版本），cuda编程入门篇见我的另一个博文哦。[借一个栗子来讲讲解基于C的cuda编程)
流程：
[*]使用pytorch训练模型，生成.pth文件；
[*]将*.pth转换成onnx模型；
[*]在tensorrt中加载onnx模型，并转换成trtde object；
[*]在trt中使用第三步转换成得object进行推理。
一、关于CenterNet后处理的思考
最终用TensorRT重构centernet，由于python代码太过于复杂，单前处理就花了很长时间。现在前处理基本上没问题，网络已经可以在tensorrt上进行forward，但是真正复杂的地方才刚开始。(注意下面的180128*128中的80代表类别，128也是相对于512的输入，down=4的来说的，如果你的输入size和类别不同，需要根据自己的实际情况来调整）
具体来说，centernet的后处理分为两个部分：
[*]对 1x80x128x128这个feature进行3x3的maxpool，其实就等于是我们传统的进行nms一样；
[*]从1x80x128x128里面根据热力图峰值点，找到我们需要的前100个目标，这里面有一个很重要的操作叫做TopK，这个topk就是将张量的最大的几个值返回，同时返回对应的index。然后根据index找到对应的中心点和偏移，对应的最终框也就找到了。
比较棘手的问题也是两个：
[*]最后的这一步maxpool和topk为什么不能放到模型里面去？而是需要后处理来算？这机加大了复杂度，同时有很浪费时间。
思考如下：这里的maxpool其实事可以放到模型里面去的，但是那样的化， 你网络出来的东西就比较多了，要出来1x80x100 同时还要对应的index。这些张量的操作放到C++里面怎么都不太好做。
主要是onnx到TensorRT的两个问题：
[*]后处理本质上是支持把这些操作同时trace到onnx模型的，它甚至不需要单独的写Plugin，因为它只需要两个onnx的op，即maxpool和topk；
[*]但问题是这两个op，在TensorRT中并没有相应的支持。
最终我们选择onnx的导出方式是，只导出网络部分，对于后处理部分，我们再单独的写layer来支持它。不过在最开始这一步是很难完成的，原因主要是：
[*]后处理有点复杂，虽然从原理上来讲，只是一个对
1x80x128x128的特征热力图的maxpool，但是自己写代码就有点繁杂了，最关键的是速度可能会很慢；
直接从网络拿到输出，再自己写代码后处理会涉及到一个topK的问题，即找到张量里面特定维度的最大值的index，这个要自己手鲁代码估计会非常非常的慢。
好在最近有朋友放出一个CUDA Kernel很好的解决了这个问题，可以说这个kernel就是本文实现CenterNet TensorRT加速的核心。这个kernel的核心逻辑，其实就是解决了上面提到的后处理的问题。这部分CUDA代码贴出来其实也十分的简单：
这部分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;
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): -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 ; c_y= grid_y + reg;
det->bbox.x1 = (c_x - wh/2)*4;
det->bbox.y1 = (c_y - wh/2)*4 ;
det->bbox.x2 = (c_x + wh/2)*4;
det->bbox.y2 = (c_y + wh/2)*4;
det->classId = cls;
det->prob = objProb;
}
}
} 需要注意的是，这里面的kernelsize就是maxpool的边长，通过设定不同的值你可以控制整个网络的粒度。最后我们通过一个threshold在这部分里面踢掉了一些置信度比较低的heappoint。
最后将结果保存为一个Detection 的数据结构中。
就这么简单，这边是CenterNet TensorRT加速的核心。最终我们也实现了通过这个后处理可以完美的得到和python一样的结果，同时拥有更快的推理速度。
未完待续
文档来源：51CTO技术博客https://blog.51cto.com/u_11495341/3036154
页:
[1]