企业官网建设流程全解析

先说一个我项目里的真实数据。同一张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 = 19200
• 40×40×3 = 4800
• 20×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.13，YOLOv5s COCO权重，conf阈值0.25，iou阈值0.45。

CPU侧慢的原因很直白。解码过滤那段要遍历25200行，每行算85个浮点数的最大值，numpy向量化能扛一部分，但argmax和条件筛选要多次扫内存；NMS就更难看了，官方实现先按score降序排，再一个个框进去和已保留框算IoU，Python循环每算一个IoU都要走一遍解释器，几百个框下来就是几毫秒。

1.3 为什么值得写核函数，而不是直接torchvision.ops.nms

这里要先说清楚适用场景，省得有人看完骂我。如果你整个pipeline都在PyTorch里跑，torchvision.ops.nms本身就是一个CUDA kernel，decode用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大小取256，grid就是(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 -std=c++17 -arch=sm_86 -Xcompiler -fPIC -shared \ yolov5_post.cu -o libyolov5_post.so

-arch要根据目标显卡选，选错了运行时大概率报“is not compatible”之类的错误，常见对应关系：

如果你要打包成库给别人用，稳妥做法是不指定-arch，改成-arch=native或者编译多个compute_XX,code=sm_XX目标再合并，代价是编译时间变长、库体积变大，但兼容性最好。

3. 核函数二：并行NMS的真面目

3.1 顺序NMS为什么在GPU上难写

传统NMS是典型的串行算法：按置信度从高到低排好序，然后从最高的框开始，逐个检查当前框和“已经保留的框”的IoU，超过阈值就扔掉。这里有个致命的数据依赖——一个框要不要被抑制，取决于它之前有哪些框被保留，而保留名单是动态累积的。这种依赖链没法直接摊到上千个线程头上，每个线程都得同步参考前面的决定，并行度起不来。

所以很多人第一版在GPU上“硬写”顺序NMS，效果都不好。正确思路是换算法，用空间换并行度。

3.2 抑制矩阵思路：每个框独立判断生死

并行NMS的核心变化是：不再维护“已保留名单”，而是把抑制关系定义成两两之间的静态关系——一个框被抑制，当且仅当存在另一个同类别、置信度严格更高（或分数相等但序号更小）的框，与它的IoU超过阈值。

这个判断只依赖全局数据，不依赖别人的判断结果，所以每个线程可以完全独立地跑完所有比较。整个流程其实退化成了：

1. 遍历所有框，找到所有“分数比我高”的框；
2. 逐个算IoU；
3. 有一个超过阈值就标记自己为抑制。

不需要排序，不需要共享内存做同步，不需要动态规划。代价是每个线程要做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 × 25200
• img_id = global_idx / 25200
• anchor = global_idx % 25200
• 输出缓冲变成[B, MAX_FILTERED, 6]，计数变成int out_count[B]，写入时用out_count + img_id

NMS仍是逐图独立的。最简单的方式是for循环发射B次NMS kernel，每次传入对应图片的偏移指针。kernel launch的开销在3~5微秒，B=8也就是40微秒，完全可接受。想再激进点可以用2D grid，一个维度是block、一个维度是图片，但代码可读性会下降，收益有限。

4.4 放进ONNX Runtime和TensorRT

如果是ONNX Runtime C++部署，GPU推理拿到的输出张量指针直接就是显存地址，ORT的GPU allocator返回的内存可以安全传给kernel。Session要配置为gpu provider，输入输出都留在GPU上，后处理完再拷最终结果。关键代码如下：

// 假设 sess 是 OrtSession，output_tensor 是 GPU 上的输出 const float* d_preds = output_tensor.GetTensorMutableData<float>(); 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的延迟变成这样：

先说清楚，绝对数值受显卡、驱动、模型、分辨率和阈值影响很大，抄作业跑不出同样数字别骂我，但相对趋势是稳定的：后处理从“和推理相当”变成“可以忽略”，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恶化的不确定项”，变成了一个可以写进延迟预算的固定值。做部署的人都知道，系统里最贵的不是慢，而是不稳定。