先说一个我项目里的真实数据。同一张640×640的图YOLOv5s在RTX 3060上推理只要5毫秒上下可后处理——就是那个没人爱看的解码加NMS——在CPU上用Python跑要7到8毫秒。推理越快后处理就越像后腿两条腿一长一短整个pipeline的耗时直接翻倍。当时我第一个念头是换CPU或者上多线程NMS但算一笔账就明白了模型输出是25200个anchor预测每个85个浮点数光把这8.5MB数据从显卡拷回内存再逐条处理就已经注定在跟PCIe带宽和CPU单核算力较劲。这篇文章记录的是我当时在这个项目里交的作业自己写CUDA核函数把YOLOv5后处理整个搬到显卡上解码、置信度过滤、NMS三段全部用kernel实现最终后处理耗时从7~8ms压到了0.2ms以内。适合正在做模型部署、被ONNX Runtime或TensorRT后处理卡住、或者想搞明白“并行NMS到底怎么整”的人看。文章会把三段流程的线程设计、完整代码、踩坑记录和实测性能数据一次讲清楚代码可以直接抄走改改用。1. 先摸清家底YOLOv5后处理到底由什么组成1.1 模型输出长什么样要写核函数第一件事是搞清楚输入数据是什么。YOLOv5在640×640输入下有三个检测头80×80、40×40、20×20每个网格点3个anchor合计anchor数量就是80×80×3 1920040×40×3 480020×20×3 1200总计 25200每个anchor是一个85维向量排布是(cx, cy, bw, bh, obj, cls_0 ... cls_79)其中cls部分按COCO数据集是80类。标准ONNX导出官方export.py出来的结果图里已经包含了sigmoid和grid偏移解码所以输出张量是[1, 25200, 85]坐标已经落在640×640的letterbox图像坐标系里objectness和类别分数也已经过sigmoid数值在0~1之间。这里的“解码”对我们来说只剩两件事第一把每个anchor的最终置信度算出来公式是score obj × max(cls)第二把xywh格式转成xyxy左上角右下角并裁剪到图像边界。有一部分部署场景导出的是raw head输出图里没有解码那核函数里就要自己补sigmoid和exp操作后面我会说明改动点。1.2 三段后处理的耗时画像我先把基线跑出来别凭感觉优化。环境是Ubuntu 20.04、RTX 3060、CUDA 11.8、PyTorch 1.13YOLOv5s COCO权重conf阈值0.25iou阈值0.45。阶段CPU numpy官方逻辑改装CUDA核函数解码 置信度过滤3.2 ms0.06 msNMS4.7 ms0.09 ms结果整理 D2H拷贝0.4 ms0.03 ms合计8.3 ms0.18 msCPU侧慢的原因很直白。解码过滤那段要遍历25200行每行算85个浮点数的最大值numpy向量化能扛一部分但argmax和条件筛选要多次扫内存NMS就更难看了官方实现先按score降序排再一个个框进去和已保留框算IoUPython循环每算一个IoU都要走一遍解释器几百个框下来就是几毫秒。1.3 为什么值得写核函数而不是直接torchvision.ops.nms这里要先说清楚适用场景省得有人看完骂我。如果你整个pipeline都在PyTorch里跑torchvision.ops.nms本身就是一个CUDA kerneldecode用tensor操作也在GPU上压根不需要自己写。我写这套核函数的真正动机是部署用ONNX Runtime或TensorRT做C部署时没有torchvision可用后处理只能在CPU上裸写模型输出8.5MB在GPU上后处理却要在CPU上做来回拷贝的D2H和H2D开销纯属浪费batch大于1时CPU侧数据量线性膨胀numpy的串行处理基本是累赘想把后处理变成一个稳定、可控、可测的自定义算子和推理算子一样纳入延迟预算。换句话说这套方案是给“需要把后处理也当成GPU算子来管理”的人准备的。纯PyTorch原型验证阶段用它属于自找麻烦。2. 核函数一解码与置信度过滤2.1 线程怎么映射到25200个anchor最简单的映射关系一个线程处理一个anchor。25200个线程block大小取256grid就是(25200 255) / 256 ≈ 99个block。每个线程的活儿完全一样读85个float算置信度够阈值就往外写。这种“同构任务”是GPU最喜欢的没有复杂分支没有线程间通信。值得说一句内存布局。preds在显存里是[25200, 85]的行优先排布线程idx读的地址是idx * 85 * 4字节起连续85个float。也就是说同一个warp里32个线程各自的起始地址相差340字节这不是教科书式的coalesced访问。但实测下来影响不大因为总共只有8.5MB数据L2缓存会把相邻cache line全兜住读一遍的耗时在60~80微秒量级。真正想压榨极限的话可以让导出端改成[85, 25200]布局再转置读但我建议别在这个阶段优化后面如果profile显示这里占大头再说。2.2 原子计数与输出缓冲设计过滤后有多少框是未知的所以不能预先知道输出长度。常规做法是开一块固定大小的显存缓冲配合atomicAdd(out_count, 1)动态分配输出槽位。我预留MAX_FILTERED 4096在多数场景下足够但如果你的conf阈值调得极低或者单张图目标特别多注意监控这个值满了会静默丢框。out_count这个计数必须每帧重置我用cudaMemsetAsync(d_filter_cnt, 0, sizeof(int), stream)放在同一个stream里保证和kernel串行执行不会和上一帧的数据混淆。2.3 完整核函数代码#include cuda_runtime.h #include math.h #define NUM_CLASSES 80 #define MAX_FILTERED 4096 // preds: [N, 85]模型输出。标准ONNX导出已包含sigmoid与grid解码 // out: [MAX_FILTERED, 6]每行存 x1, y1, x2, y2, score, label __global__ void decode_filter_kernel( const float* preds, float* out, int* out_count, int N, float conf_threshold, float img_w, float img_h) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx N) return; const float* p preds idx * (5 NUM_CLASSES); // 如果你的导出图没做sigmoid这里要补上 // float cx 1.f / (1.f expf(-p[0])); // float cy 1.f / (1.f expf(-p[1])); // float obj 1.f / (1.f expf(-p[4])); // w/h 要改成 expf(p[2]) / expf(p[3]) float cx p[0]; float cy p[1]; float bw p[2]; float bh p[3]; float obj p[4]; int best_cls 0; float max_cls p[5]; for (int c 1; c NUM_CLASSES; c) { if (p[5 c] max_cls) { max_cls p[5 c]; best_cls c; } } float score obj * max_cls; if (score conf_threshold) return; float x1 fmaxf(0.f, cx - bw * 0.5f); float y1 fmaxf(0.f, cy - bh * 0.5f); float x2 fminf(img_w, cx bw * 0.5f); float y2 fminf(img_h, cy bh * 0.5f); int pos atomicAdd(out_count, 1); if (pos MAX_FILTERED) return; float* o out pos * 6; o[0] x1; o[1] y1; o[2] x2; o[3] y2; o[4] score; o[5] (float)best_cls; }编译命令很简单nvcc -O3 -stdc17 -archsm_86 -Xcompiler -fPIC -shared \ yolov5_post.cu -o libyolov5_post.so-arch要根据目标显卡选选错了运行时大概率报“is not compatible”之类的错误常见对应关系算力代号典型显卡sm_75T4、RTX 20系列sm_80A100sm_86RTX 30系列sm_89RTX 40系列sm_120RTX 50系列如果你要打包成库给别人用稳妥做法是不指定-arch改成-archnative或者编译多个compute_XX,codesm_XX目标再合并代价是编译时间变长、库体积变大但兼容性最好。3. 核函数二并行NMS的真面目3.1 顺序NMS为什么在GPU上难写传统NMS是典型的串行算法按置信度从高到低排好序然后从最高的框开始逐个检查当前框和“已经保留的框”的IoU超过阈值就扔掉。这里有个致命的数据依赖——一个框要不要被抑制取决于它之前有哪些框被保留而保留名单是动态累积的。这种依赖链没法直接摊到上千个线程头上每个线程都得同步参考前面的决定并行度起不来。所以很多人第一版在GPU上“硬写”顺序NMS效果都不好。正确思路是换算法用空间换并行度。3.2 抑制矩阵思路每个框独立判断生死并行NMS的核心变化是不再维护“已保留名单”而是把抑制关系定义成两两之间的静态关系——一个框被抑制当且仅当存在另一个同类别、置信度严格更高或分数相等但序号更小的框与它的IoU超过阈值。这个判断只依赖全局数据不依赖别人的判断结果所以每个线程可以完全独立地跑完所有比较。整个流程其实退化成了遍历所有框找到所有“分数比我高”的框逐个算IoU有一个超过阈值就标记自己为抑制。不需要排序不需要共享内存做同步不需要动态规划。代价是每个线程要做O(M²)次遍历但M是过滤后的候选框数量通常只有几百个几千次浮点运算对GPU来说眨眼就完。注意一个细节YOLOv5的NMS是按类别独立的两个不同类别的框即使完全重叠也应该都保留。所以核函数里比较之前必须先判断label相等这个最容易漏。3.3 完整NMS核函数代码// boxes: [M, 6]就是过滤kernel的输出 // m_dev: GPU上的过滤计数由上一个kernel写入避免host同步 // keep: [M] 输出0/1 __global__ void parallel_nms_kernel( const float* boxes, const int* m_dev, unsigned char* keep, float iou_threshold) { int M min(MAX_FILTERED, *m_dev); int i blockIdx.x * blockDim.x threadIdx.x; if (i M) return; const float* bi boxes i * 6; float x1 bi[0], y1 bi[1], x2 bi[2], y2 bi[3]; float score_i bi[4]; int label_i (int)bi[5]; float area_i (x2 - x1) * (y2 - y1); bool suppressed false; for (int j 0; j M; j) { if (j i) continue; const float* bj boxes j * 6; int label_j (int)bj[5]; if (label_j ! label_i) continue; float score_j bj[4]; // 分数更高或同级但序号更小的框才有资格抑制当前框 if (score_j score_i || (score_j score_i j i)) continue; float inter_w fmaxf(0.f, fminf(x2, bj[2]) - fmaxf(x1, bj[0])); float inter_h fmaxf(0.f, fminf(y2, bj[3]) - fmaxf(y1, bj[1])); float inter inter_w * inter_h; float area_j (bj[2] - bj[0]) * (bj[3] - bj[1]); float iou inter / (area_i area_j - inter 1e-6f); if (iou iou_threshold) { suppressed true; break; } } keep[i] suppressed ? 0 : 1; }NMS跑完keep数组里还有0/1标记但输出框要求是紧凑排列的所以还需要一个打包kernel把保留的框连续写进最终结果__global__ void compact_kernel( const float* boxes, const unsigned char* keep, const int* m_dev, float* out, int* out_count) { int M min(MAX_FILTERED, *m_dev); int i blockIdx.x * blockDim.x threadIdx.x; if (i M) return; if (!keep[i]) return; int pos atomicAdd(out_count, 1); if (pos MAX_FILTERED) return; const float* src boxes i * 6; float* dst out pos * 6; #pragma unroll for (int k 0; k 6; k) dst[k] src[k]; }3.4 并行NMS和顺序NMS的差异别装看不见必须承认并行NMS的结果和官方顺序NMS并不严格等价。举一个例子A、B、C三个框分数A最高、B其次、C最低A和B重叠严重B和C重叠严重但A和C不重叠。顺序NMS里A保留B被A抑制C再进来时只跟保留名单比只跟A比不重叠所以C保留。并行NMS里C会同时跟A和B比发现B比它分高且重叠于是C被抑制。这种“被抑制的框又去抑制别人”的现象在并行NMS里是存在的。在真实检测场景里同类目标连成一串且两两重叠的场景很少见我拿COCO val跑过一遍mAP差异在0.1%以内基本是噪声级别。但如果你下游要严格对齐PyTorch官方输出必须在测试里盯这一点。真想完全等价得回到“只跟保留框比”的老路那就得牺牲并行度性价比不高我的建议是接受差异并写好测试用例兜底。4. 工程化落地内存复用、batch处理与推理框架衔接4.1 显存缓冲只分配一次写核函数最忌讳每帧cudaMalloc。显存分配走驱动开销在微秒到几十微秒波动而且会跟CUDA context的锁纠缠延迟不稳定。正确姿势是静态或全局指针第一次调用时分配之后所有帧复用。上面代码里我用static指针做懒分配实际工程里可以封装成类在构造函数里分配析构里释放。每帧要做的只有两个cudaMemsetAsync重置过滤计数和最终计数这两个操作都在stream里异步执行不会阻塞。4.2 怎么把M从设备端拿到主机端而不卡流水线这是后处理工程化里最容易踩的坑。NMS核函数和compact核函数都要知道过滤后的框数M但M在设备端。如果每帧都cudaMemcpyAsync加cudaStreamSynchronize等主机读到M再发下一个kernel流水线就直接被sync打断了延迟全回来了。我用的方案是设备端直接读计数NMS和compact核函数内部通过min(MAX_FILTERED, *m_dev)拿到M主机端完全不需要知道中间值。三个kernel用固定grid尺寸连续发射中间零同步。真正的同步只剩最后一步把最终结果从设备拷回主机。最终结果拷贝也有讲究。虽然compact后框数写在了d_final_cnt里但你不想先拷计数再拷数据那就按固定大小拷直接拷MAX_FILTERED * 6 * sizeof(float)48KB的量对PCIe来说不到30微秒拷回来之后只读前count个框就行。省了同步多了点带宽划算。如果你对延迟极致敏感还有第三招用cudaMemcpyAsync拷到pinned memory然后主机侧轮询pinned内存里的计数变化这就是另一个话题了。4.3 batch怎么处理batch处理其实不复杂核心是给每张图独立的计数。过滤kernel把一维索引拆成图和anchor两个维度总线程数 B × 25200img_id global_idx / 25200anchor global_idx % 25200输出缓冲变成[B, MAX_FILTERED, 6]计数变成int out_count[B]写入时用out_count img_idNMS仍是逐图独立的。最简单的方式是for循环发射B次NMS kernel每次传入对应图片的偏移指针。kernel launch的开销在3~5微秒B8也就是40微秒完全可接受。想再激进点可以用2D grid一个维度是block、一个维度是图片但代码可读性会下降收益有限。4.4 放进ONNX Runtime和TensorRT如果是ONNX Runtime C部署GPU推理拿到的输出张量指针直接就是显存地址ORT的GPU allocator返回的内存可以安全传给kernel。Session要配置为gpu provider输入输出都留在GPU上后处理完再拷最终结果。关键代码如下// 假设 sess 是 OrtSessionoutput_tensor 是 GPU 上的输出 const float* d_preds output_tensor.GetTensorMutableDatafloat(); yolov5_postprocess(d_preds, 25200, d_final, d_final_count, conf_th, iou_th, img_w, img_h, stream); int count 0; cudaMemcpy(count, d_final_count, sizeof(int), cudaMemcpyDeviceToHost); cudaMemcpy(h_final, d_final, count * 6 * sizeof(float), cudaMemcpyDeviceToHost);TensorRT的话最规矩的做法是做成plugin把这三个kernel串进plugin的enqueue里这样TensorRT就能把后处理和前向推理放进同一个CUDA graph延迟最稳。如果不想写plugin也可以直接在推理完成后读输出再调这套函数效果略差但省事。5. 实测数据与调参避坑5.1 端到端对比同样的机器、同样的模型、同样的阈值把后处理从CPU换成CUDA核函数之后整条pipeline的延迟变成这样场景推理(GPU)后处理端到端原方案 batch14.8 ms8.3 msCPU13.1 ms本方案 batch14.8 ms0.18 ms4.98 ms原方案 batch816.2 ms约45 msCPU串行61 ms本方案 batch816.2 ms0.9 ms17.1 ms先说清楚绝对数值受显卡、驱动、模型、分辨率和阈值影响很大抄作业跑不出同样数字别骂我但相对趋势是稳定的后处理从“和推理相当”变成“可以忽略”batch越大收益越明显。5.2 我踩过的几个坑第一个坑是忘了重置计数。第一版跑出来的框数每帧递增因为out_count没清零原子加一直在旧值基础上累加。排查过程很痛苦因为框的排列是乱的看起来像是随机多框。后来在Nsight里看到计数内存的读写轨迹才定位。从那以后我把cudaMemsetAsync放在了三个kernel的最前面并统一在同一个stream里保证顺序。第二个坑是输出框顺序不稳定。atomicAdd分配的槽位是乱序的最终结果的框不是按置信度从高到低排的。官方NMS的排序结果被下游依赖时这里就会出问题。我的处理是在最后一个compact之后对拷回主机的几十个框用std::sort按分数排一下量小成本几乎为零。第三个坑是letterbox坐标还原。YOLOv5预处理是等比缩放加灰边填充模型输出坐标在640×640的letterbox坐标系里。直接拿去画框会整体偏移。正确做法是核函数里只管裁剪到640边界拷回主机后对最终少量框统一做x (x - pad_x) / ratio的还原。别在核函数里逐anchor做除法白白增加计算量。第四个坑是float16输入。TensorRT开FP16之后模型输出指针指向的是__half数组直接用float*读会读到一堆乱码。要么在插件里加一个__half*版本的过滤核函数要么先把FP16输出转成FP32再进后处理。转换本身也就是一次elementwise kernel的事情但忘记处理就很致命。5.3 调参经验block size我用256对于过滤kernel和compact kernel都合适。NMS核函数里每个线程要遍历全部M个框block大小对性能影响不大但注意MAX_FILTERED别开太大4096已经覆盖绝大多数情况。如果你把conf阈值降到0.05还想保持全量框建议先把过滤逻辑改成“每个block先局部筛一遍再合并”避免大量线程同时去抢一个原子计数器。还有一个容易忽略的点CUDA context初始化。第一次调用任何CUDA kernel都有上下文创建和模块加载的开销几百毫秒级别。做性能测试或者上线前预热至少空跑两三帧再开始计时否则会看到一个吓人的首帧延迟。最后补一句体感把这三个kernel串进同一个stream之后我打开Nsight看时间轴整个后处理段变成了一条不到0.3毫秒的细线和旁边CPU侧那段忽高忽低的红色长条形成鲜明对比。对我来说这件事最大的收获不是省了多少毫秒而是后处理的耗时从“随框数波动、随batch恶化的不确定项”变成了一个可以写进延迟预算的固定值。做部署的人都知道系统里最贵的不是慢而是不稳定。