torchvision.ops.nms默认跑在cpu上是因为其贪心迭代逻辑（选最高分框→逐个计算iou→删除超阈值框）天然依赖串行执行，早期实现复用cpu版本且未强制设备检查，即使输入cuda张量也会悄悄回拷cpu执行；这是历史兼容性选择，非bug。
PyTorch的torchvision.ops.nms为什么默认跑在CPU上？
因为它的实现逻辑是贪心迭代：每次选最高分框，再逐个算它和剩下所有框的IoU，删掉超阈值的——这种“选一个、算一批、删一批”的依赖链，天然不适合GPU并行。PyTorch早期直接复用CPU版逻辑，没强制要求输入在GPU上，所以哪怕你传的是cuda张量，内部也会悄悄拷回CPU执行。这不是bug，是历史兼容性选择。
什么时候必须把NMS搬到GPU上？
当你的检测模型每帧输出超过5000个候选框，且帧率要求≥30fps时，CPU NMS就会成为瓶颈。典型场景包括：
- YOLOv5/v8在640×640输入下原始输出约25200个框
- 工业质检相机120fps采集，NMS耗时超过3ms就会拖垮整条流水线
- 多尺度特征融合后，不同anchor层输出叠加，框数翻倍
这时torchvision.ops.nms若仍在CPU跑，光是25200个框从GPU→CPU→GPU的拷贝，单帧就吃掉400KB PCIe带宽，实测延迟飙升至8–12ms。
torchvision.ops.nms在GPU上运行的硬性条件
三个条件缺一不可：
- 输入
boxes和scores必须是cuda设备上的张量（不能是cpu或mps） - PyTorch版本 ≥ 1.10（旧版本即使传cuda张量也fallback到CPU）
- 显卡驱动和CUDA toolkit匹配（常见坑：CUDA 11.8驱动装了但PyTorch编译用的是11.3，NMS内核加载失败，静默退回到CPU）
验证方式很简单：torch.cuda.current_device()之后立刻调用nms，用nvidia-smi看GPU显存占用是否跳变、计算时间是否压到0.5ms以内。
为什么不用自己手写CUDA NMS？
除非你真要压榨最后0.1ms——比如Jetson Orin部署时目标是单帧≤4ms。否则torchvision.ops.nms已足够：它底层调用的是Facebook优化过的C++/CUDA混合实现，支持batched NMS、自动内存复用，且与ONNX导出兼容。自己手写容易踩的坑包括：
- IoU计算里漏掉
torch.clamp(..., min=0)，导致负面积除零崩溃 - 排序索引没用
argsort(descending=True)，结果顺序错乱 - 没处理
boxes坐标越界（x1≥x2或y1≥y2），GPU kernel直接abort
真正该花时间调的，是iou_threshold和前置的置信度阈值——这两个参数对mAP的影响，远大于换用自定义NMS带来的那0.3ms收益。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
