目录
前言
减少YOLO后处理nms的计算和比较次数。
YOLO-det输出的维度是(1, 4+cls_num, 8400),如果直接进行nms比较会进行大量无效的重复操作:排序中的比较、四则运算等。
1.问题分析
nms的python代码:
def nms(boxes, scores, iou_threshold):
boxes = np.array(boxes)
scores = np.array(scores)
# Sort by score
sorted_indices = np.argsort(scores)[::-1]
keep_boxes = []
while sorted_indices.size > 0:
# Pick the last box
box_id = sorted_indices[0]
keep_boxes.append(box_id)
# Compute IoU of the picked box with the rest
ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
# Remove boxes with IoU over the threshold
keep_indices = np.where(ious < iou_threshold)[0]
# print(keep_indices.shape, sorted_indices.shape)
sorted_indices = sorted_indices[keep_indices + 1]
return keep_boxes
问题1:排序问题
上述代码先对所有检测结果进行排序,然后再计算IoU,lg8400≈13,则实际排序中的比较次数nlgn约为8400×13≈110,000次。
问题2:极大值抑制问题
极大值抑制需要两两比较,时间复杂度为O(n^2),这个就更高了,不过由于已经使用置信度筛选了,一般也就计算个百来次。
2.优化比较和计算次数
C++的nms代码照着Python实现,不需要改变,修改传入前的bbox:
int num_channels = 5;
int num_detections = output_size / num_channels;
// 定义置信度阈值
float confidence_threshold = 0.5f; // 根据需要调整
// 提取满足置信度阈值的检测框和置信度,同时进行xywh比较优化
std::vector<std::vector<float>> filtered_boxes;
std::vector<float> confidences;
for (int i = 0; i < num_detections; ++i) {
float cx = output_data[i]; // 中心点x
float cy = output_data[num_detections + i]; // 中心点y
float w = output_data[2 * num_detections + i]; // 宽度
float h = output_data[3 * num_detections + i]; // 高度
float confidence = output_data[4 * num_detections + i]; // 置信度
if (confidence > confidence_threshold) {
if (!filtered_boxes.empty()) {
// 获取前一个满足条件的检测框
const std::vector<float>& last_box = filtered_boxes.back();
// 比较当前检测框与前一个检测框的xywh差值
if (std::fabs(cx - last_box[0]) < 1.0f &&
std::fabs(cy - last_box[1]) < 1.0f &&
std::fabs(w - last_box[2]) < 1.0f &&
std::fabs(h - last_box[3]) < 1.0f) {
// 如果当前置信度大于上一个,则替换上一个检测框
if (confidence > confidences.back()) {
filtered_boxes.back() = {cx, cy, w, h};
confidences.back() = confidence;
}
continue; // 当前值被舍弃或替换,跳过后续步骤
}
}
// 如果没有相似的前一个检测框,或当前框不被舍弃,则保存当前检测框
filtered_boxes.push_back({cx, cy, w, h});
confidences.push_back(confidence);
}
}
// 将xywh转换为xyxy
std::vector<std::vector<float>> boxes;
boxes.reserve(filtered_boxes.size());
for (const auto& box : filtered_boxes) {
float cx = box[0];
float cy = box[1];
float w = box[2];
float h = box[3];
float x1 = cx - w / 2.0f;
float y1 = cy - h / 2.0f;
float x2 = cx + w / 2.0f;
float y2 = cy + h / 2.0f;
boxes.push_back({x1, y1, x2, y2});
}
std::cout << "Boxes size before NMS: " << boxes.size() << std::endl;
优化1:跳过reshape直接置信度筛选
原始Python会把推理输出的维度(1, 4+cls_num, 8400) reshape成(8400, 4+cls_num),再进行后续操作(所有操作是对原本所有预测结果进行的)。
C++中多维数组本身就是一维数组,所以直接单循环遍历一次即可,在置信度满足时,将值保存到新的变量filtered_boxes中(如果为了省去开辟空间,可以将值保存到原vector的前面,并记录末尾索引):
// 提取满足置信度阈值的检测框和置信度,同时进行xywh比较优化
std::vector<std::vector<float>> filtered_boxes;
std::vector<float> confidences;
for (int i = 0; i < num_detections; ++i) {
float cx = output_data[i]; // 中心点x
float cy = output_data[num_detections + i]; // 中心点y
float w = output_data[2 * num_detections + i]; // 宽度
float h = output_data[3 * num_detections + i]; // 高度
float confidence = output_data[4 * num_detections + i]; // 置信度
if (confidence > confidence_threshold) {
// 置信度时满足操作
}
优化2:减少用于nms的bbox数
经过优化1后,用于nms的bbox共91个,可以观察到,由于推理是“滑动窗口式”地进行的,会使得“重复”预测的值在“排列”上是相邻的,且重复预测的值在xywh上差值往往在“1.0”以内:
因此,可以进行一次遍历,用当前值和最后一次满足条件的值对比,如果四个差值均在“1.0”以内,则保留较大置信度的值:
if (!filtered_boxes.empty()) {
// 获取前一个满足条件的检测框
const std::vector<float>& last_box = filtered_boxes.back();
// 比较当前检测框与前一个检测框的xywh差值
if (std::fabs(cx - last_box[0]) < 1.0f &&
std::fabs(cy - last_box[1]) < 1.0f &&
std::fabs(w - last_box[2]) < 1.0f &&
std::fabs(h - last_box[3]) < 1.0f) {
// 如果当前置信度大于上一个,则替换上一个检测框
if (confidence > confidences.back()) {
filtered_boxes.back() = {cx, cy, w, h};
confidences.back() = confidence;
}
continue; // 当前值被舍弃或替换,跳过后续步骤
}
}
上述处理后,用于nms的bbox由91降至34:
原始IoU比较次数:91×91÷2=4140次;
修改后IoU比较次数:34×34÷2=578次+91次xywh比较。减少了86%的IoU比较次数。(xywh比较仅涉及4次减法,运算量远远小于IoU比较)。
算法有些误差(原始nms计算会给无效的值打上mask,计算量会减少不少),但也能直观反应比较次数的减少。
3.举个荔枝
修改前,一个目标有10个bbox,进行了10次IoU比较:
修改后,一个目标仅有1个bbox,IoU比较一次都没进行:
总结
优化1:将110,000次的比较优化到9,000次(Python实现和C++实现的对比,网上找的代码,如果是官方一点的,应该也有类似优化)。
优化2:将O(NlgN) 的计算复杂度(排序比较+IoU比较)降到接近O(N)(实际效果可能更好,当然不排除特别极端情况下,也会变差,不过这是基于YOLO特性设计的,几乎不可能出现)。