TensorRT+CUDA加速优化版CenterNet旋转目标以及水平目标框的检测

2024-01-29 15:47:04

前言

由于工作项目所需，一直用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[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后处理的思考

相关文章