MXNet的高效率CUDA NMS解析

最新推荐文章于 2024-10-13 11:35:14 发布

Law-Yao

最新推荐文章于 2024-10-13 11:35:14 发布

阅读量1.8k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： CUDA 优化加速深度学习文章标签： CUDA

本文链接：https://blog.youkuaiyun.com/nature553863/article/details/91813017

非极大抑制（NMS）在目标检测等领域用于消除重叠边框，MXNet的MultiBoxDetection提供了高效的CUDA实现。通过mx.contrib.ndarray.MultiBoxDetection API，结合CUDA代码解析，能显著提升计算速度。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

非极大抑制（NMS：Non Maximum Suppression）起到边框（水平框或倾斜框）去重叠的作用，广泛应用于通用目标检测、人脸检测与OCR检测等算法的后处理。简单的Python实现如下：

def standard_nms(S, thres):
    order = np.argsort(S[:, -1])[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        ovr = np.array([intersection(S[i], S[t]) for t in order[1:]])

        inds = np.where(ovr <= thres)[0]
        order = order[inds+1]
    return S[keep]

MXNet的MultiBoxDetection涉及的NMS，体现了非常高的计算效率，API调用接口: mx.contrib.ndarray.MultiBoxDetection，CUDA代码链接以及具体解析如下：

https://github.com/apache/incubator-mxnet/blob/master/src/operator/contrib/multibox_detection.cu

// 限定数值取值范围；__device__声明的核函数只能被CUDA核函数调用
template<typename DType>
__device__ void Clip(DType *value, const DType lower, const DType upper) {
  if ((*value) < lower) *value = lower;
  if ((*value) > upper) *value = upper;
}

// 计算两个水平框的IOU=交集面积/并集面积
template<typename DType>
__device__ void CalculateOverlap(const DType *a, const DType *b, DType *iou) {
  DType w = max(DType(0), min(a[2], b[2]) - max(a[0], b[0]));
  DType h = max(DType(0), min(a[3], b[3]) - max(a[1], b[1]));
  DType i = w * h;
  DType u = (a[2] - a[0]) * (