C++部署YOLO算法实战：从模型转换到性能优化全解析

1次阅读

共计 3110 个字符，预计需要花费 8 分钟才能阅读完成。

最近在做一个智能监控项目，需要把 YOLOv5 模型部署到 C ++ 生产环境。本以为直接调用 PyTorch 的 LibTorch 就完事了，结果踩坑踩到怀疑人生。这里总结几个典型问题：

模型格式问题：PyTorch 的.pt 模型不能直接被 C ++ 加载，需要转换中间格式
依赖地狱：LibTorch 的 ABI 兼容性问题导致编译时各种符号冲突
性能瓶颈：原生 Python 前处理 +OpenCV 混合调用导致 GPU 利用率不足 50%

花了两周对比主流方案，测试数据供参考：

引擎	推理速度(ms)	内存占用(MB)	易用性
ONNX Runtime	45	1200	★★★★
TensorRT	22	800	★★
LibTorch	60	1500	★★★

选型建议：
– 快速验证选 ONNX Runtime
– 极致性能选 TensorRT
– 需要动态输入选 LibTorch

用官方 export.py 脚本时要注意两个关键参数：

# 导出动态 batch 模型（必须设置！）torch.onnx.export(
    model,
    im,
    'yolov5s.onnx',
    opset_version=12,
    input_names=['images'],
    output_names=['output'],
    dynamic_axes={'images': {0: 'batch'},  # 动态 batch 维度
        'output': {0: 'batch'}
    })

用 OpenCV 做图像预处理时，注意 BGR 到 RGB 的转换：

cv::Mat preprocess(cv::Mat& img) {
    // 保持长宽比 resize
    float scale = std::min(640.0f/img.cols, 640.0f/img.rows);
    cv::Mat resized;
    cv::resize(img, resized, cv::Size(), scale, scale);

    // 制作 640x640 画布
    cv::Mat canvas = cv::Mat::zeros(640, 640, CV_8UC3);
    resized.copyTo(canvas(cv::Rect(0, 0, resized.cols, resized.rows)));

    // 归一化并转为 NCHW 格式
    cv::dnn::blobFromImage(canvas, tensor, 1/255.0, 
                          cv::Size(640,640), 
                          cv::Scalar(0,0,0), true, false);
    return tensor;
}

YOLO 的输出解码和 NMS 是最耗时的部分，这里给出 OpenMP 加速版：

void postprocess(float* output, std::vector<Detection>& dets) {
    #pragma omp parallel for
    for (int i = 0; i < num_anchors; ++i) {
        // 解析 box 坐标
        float* ptr = output + i * 85;
        float conf = ptr[4];
        if (conf < threshold) continue;

        // 并行处理每个 anchor
        Detection det;
        // ... 解码逻辑
        #pragma omp critical
        dets.emplace_back(det);
    }

    // 快速 NMS 实现
    std::sort(dets.begin(), dets.end(), 
             [](auto& a, auto& b) {return a.conf > b.conf;});

    for (size_t i = 0; i < dets.size(); ++i) {if (dets[i].conf == 0) continue;
        for (size_t j = i + 1; j < dets.size(); ++j) {if (iou(dets[i], dets[j]) > nms_thresh)
                dets[j].conf = 0;
        }
    }
}

用 trtexec 工具转换时，开启 FP16 模式能直接提升 2 倍速度：

trtexec --onnx=yolov5s.onnx \
        --saveEngine=yolov5s_fp16.engine \
        --fp16 \
        --workspace=2048

建议采用生产者 - 消费者模式：

class InferQueue {
public:
    void start() {workers_.reserve(num_threads);
        for(int i=0; i<num_threads; ++i) {workers_.emplace_back([this]{while(!stop_) {
                    Task task;
                    if(queue_.try_pop(task)) {auto result = engine_->infer(task.input);
                        task.promise.set_value(result);
                    }
                }
            });
        }
    }

    std::future<Result> submit(cv::Mat input) {
        std::promise<Result> prom;
        auto fut = prom.get_future();
        queue_.push({input, std::move(prom)});
        return fut;
    }

private:
    moodycamel::ConcurrentQueue<Task> queue_; // 无锁队列
    std::vector<std::thread> workers_;
};

反复申请释放显存会导致性能下降，建议预分配：

class MemoryPool {
public:
    void* allocate(size_t size) {std::lock_guard<std::mutex> lock(mutex_);
        auto it = std::find_if(pool_.begin(), pool_.end(), 
                              [size](auto& chunk){return !chunk.used && chunk.size >= size;});
        if (it != pool_.end()) {
            it->used = true;
            return it->ptr;
        }

        // 没有可用块就新建
        void* ptr = cudaMalloc(size);
        pool_.push_back({ptr, size, true});
        return ptr;
    }

private:
    struct MemChunk {
        void* ptr;
        size_t size;
        bool used;
    };
    std::vector<MemChunk> pool_;
    std::mutex mutex_;
};