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

1次阅读

没有评论

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

在实际项目中，很多团队会选择 PyTorch 训练 YOLO 模型，但生产环境往往需要 C ++ 部署。这个过程中会遇到几个典型问题：

PyTorch 模型直接导出为 TorchScript 后，在 C ++ 端依赖完整的 LibTorch 库，体积庞大（超过 1GB）
不同版本的 CUDA/cuDNN 与推理引擎存在兼容性问题，经常出现版本冲突
原生 Python 实现的 NMS 等后处理逻辑需要重写为 C ++，且要保证数值计算一致性

目前主流的 C ++ 推理方案有三种：

ONNX Runtime
优点：跨平台支持好，API 简单，支持动态输入
缺点：性能中等，量化支持有限
TensorRT
优点：极致性能，支持多种量化（FP16/INT8）
缺点：生态封闭，模型转换复杂
LibTorch
优点：与训练框架无缝对接
缺点：依赖庞大，性能较差

对于大多数场景，推荐组合方案：

flowchart LR
    PyTorch 训练 --> ONNX 导出 --> TensorRT 优化 --> C++ 部署

使用官方 export.py 脚本时需特别注意动态维度：

# 关键导出参数
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 进行高效图像预处理：

// 现代 C ++ 风格的前处理
cv::Mat preprocess(cv::Mat& img, int target_size) {
    cv::Mat resized;
    cv::resize(img, resized, cv::Size(target_size, target_size));

    // 归一化操作 (0- 1 范围)
    resized.convertTo(resized, CV_32F, 1.0/255);

    // HWC 转 CHW
    cv::dnn::blobFromImage(resized, resized);
    return resized;
}

手写 CUDA 核函数加速 NMS：

__global__ void nms_kernel(float* boxes, float* scores, ...) {
    // 共享内存加速 IO
    __shared__ float local_boxes[THREADS_PER_BLOCK * 6];
    ...
}

构建器配置关键参数：

config->setFlag(BuilderFlag::kFP16);
config->setMaxWorkspaceSize(1 << 30);

class MemoryPool {
    std::vector<void*> gpu_ptrs_;
    cudaStream_t stream_;
public:
    void* allocate(size_t size) {
        void* ptr;
        cudaMallocAsync(&ptr, size, stream_);
        gpu_ptrs_.push_back(ptr);
        return ptr;
    }
    ~MemoryPool() {for(auto ptr : gpu_ptrs_) 
            cudaFreeAsync(ptr, stream_);
    }
};