CANN/cann-bench NMS算子API描述

NMS 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介对候选框执行非极大值抑制 (Non-Maximum Suppression)根据置信度分数和 IoU 阈值过滤重叠的候选框。主要应用场景目标检测模型如 YOLO、Faster R-CNN的后处理去重实例分割中候选区域的筛选人脸检测等需要消除重叠检测框的场景算子特征难度等级L3SortSelect双输入单输出根据置信度排序并基于 IoU 阈值迭代过滤重叠框输出保留框的索引2. 算子定义数学公式$$ \text{keep_indices} \text{NMS}(\text{boxes}, \text{scores}, \text{iou_threshold}) $$处理流程按 scores 从高到低对候选框排序选取得分最高的框加入保留列表计算该框与剩余所有框的 IoU交并比移除 IoU 大于iou_threshold的重叠框重复步骤 2-4直至所有框被处理其中 IoU 定义为$$ \text{IoU}(A, B) \frac{|A \cap B|}{|A \cup B|} $$3. 接口规范算子原型cann_bench.nms(Tensor boxes, Tensor scores, float iou_threshold) - Tensor keep_indices输入参数说明参数类型默认值描述boxesTensor必选输入候选框格式为 [x1, y1, x2, y2]shape 为 [N, 4]scoresTensor必选每个候选框的置信度分数shape 为 [N]iou_thresholdfloat必选IoU 阈值用于过滤重叠框输出参数Shapedtype描述keep_indices[M]int64NMS 后保留的框索引M N数据类型输入 (boxes) dtype输入 (scores) dtype输出 dtypefloat32float32int64规则与约束boxes的形状必须为 [N, 4]每行为 [x1, y1, x2, y2] 格式scores的形状必须为 [N]且 N 与 boxes 的第一维一致iou_threshold取值范围为 (0, 1)值越小过滤越严格输出keep_indices为 1D int64 张量长度 M 取决于过滤后保留的框数输出索引按置信度从高到低排序支持范围输入 tensor 各维度与参数的支持范围维度 / 参数范围备注boxesshape[N, 4]N ∈ [1, 8192]cases.csv 实测 N ∈ [511, 4096]包含对齐 (1024/4096) 与质数非对齐 (1009/3001/4001)scoresshape[N]与 boxes 第一维一致cases.csv 实测 N ∈ [511, 4096]boxesdtypefloat32cases.csv 实测仅 float32scoresdtypefloat32cases.csv 实测仅 float32boxesvaluefloat32 有限值范围cases.csv 实测 [-65504, 65504]含 float16 边界 / ±inf / nan / 全零等特殊值scoresvalue[0, 1]cases.csv 实测均为 [0, 1]置信度语义iou_threshold(0, 1)cases.csv 实测 0.05 ~ 0.9覆盖严格 / 宽松阈值输出keep_indices[M]M ≤ Nint64长度由过滤结果决定按 scores 从高到低排序约束scores.shape[0]必须等于boxes.shape[0]iou_threshold须严格落在开区间 (0, 1) 内。4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch NMS 算子 Torch Golden 参考实现 对候选框执行非极大值抑制 (Non-Maximum Suppression) 公式keep_indices nms(boxes, scores, iou_threshold) def nms( boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float ) - torch.Tensor: 对候选框执行非极大值抑制 公式keep_indices nms(boxes, scores, iou_threshold) Args: boxes: 输入候选框格式为 [x1, y1, x2, y2]shape 为 [N, 4] scores: 每个候选框的置信度分数shape 为 [N] iou_threshold: IoU 阈值用于过滤重叠框 Returns: keep_indices: NMS 后保留的框索引shape 为 [M] # 确保输入格式正确 assert boxes.dim() 2 and boxes.shape[1] 4, boxes shape must be [N, 4] assert scores.dim() 1 and scores.shape[0] boxes.shape[0], scores shape must be [N] # 纯 PyTorch 实现 NMS避免 torchvision ABI 兼容问题 areas (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) _, order scores.sort(descendingTrue) keep [] while order.numel() 0: if order.numel() 0: break i order[0].item() keep.append(i) if order.numel() 1: order order.new_empty(0) break xx1 boxes[order[1:], 0].clamp(minboxes[i, 0]) yy1 boxes[order[1:], 1].clamp(minboxes[i, 1]) xx2 boxes[order[1:], 2].clamp(maxboxes[i, 2]) yy2 boxes[order[1:], 3].clamp(maxboxes[i, 3]) w (xx2 - xx1).clamp(min0) h (yy2 - yy1).clamp(min0) inter w * h iou inter / (areas[i] areas[order[1:]] - inter 1e-6) inds (iou iou_threshold).nonzero(as_tupleFalse).squeeze(1) order order[inds 1] return torch.tensor(keep, dtypetorch.long, deviceboxes.device)6. 额外信息算子调用示例import torch import cann_bench boxes torch.rand(1000, 4, dtypetorch.float32, devicenpu) * 100 scores torch.rand(1000, dtypetorch.float32, devicenpu) keep cann_bench.nms(boxes, scores, iou_threshold0.5) # 低 IoU 阈值更严格的过滤 keep cann_bench.nms(boxes, scores, iou_threshold0.3)【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考