c++如何實現yolov5轉onnx

# C++如何實現YOLOv5轉ONNX

## 前言

在深度學習模型部署過程中，模型格式轉換是至關重要的一環。YOLOv5作為當前最流行的目標檢測算法之一，其從PyTorch到ONNX的轉換過程需要特別注意細節。本文將深入探討如何使用C++環境實現YOLOv5模型到ONNX格式的轉換，涵蓋從環境準備到最終導出的完整流程。

## 一、環境準備

### 1.1 基礎軟件依賴

```bash
# 必需組件清單
- CMake 3.12+
- GCC 7.0+/Clang 5.0+
- Python 3.7+ (用于運行轉換腳本)
- PyTorch 1.7.0+
- LibTorch (C++版PyTorch)
- ONNX Runtime 1.8+

1.2 開發環境配置

// 示例：檢查環境版本的C++代碼
#include <iostream>
#include <torch/version.h>
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>

int main() {
    std::cout << "LibTorch版本: " << TORCH_VERSION << std::endl;
    Ort::Env env;
    std::cout << "ONNX Runtime版本: " << env.GetVersionString() << std::endl;
    return 0;
}

1.3 YOLOv5模型獲取

建議直接從官方倉庫克隆最新代碼：

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

二、PyTorch模型導出基礎

2.1 原生PyTorch導出方法

YOLOv5官方提供的export.py腳本支持多種格式導出：

# 基本導出命令
python export.py --weights yolov5s.pt --include onnx

2.2 導出參數詳解

三、C++實現轉換的核心邏輯

3.1 使用LibTorch加載模型

#include <torch/script.h>

torch::jit::script::Module load_model(const std::string& model_path) {
    try {
        // 加載PyTorch模型
        auto module = torch::jit::load(model_path);
        module.eval();
        return module;
    } catch (const c10::Error& e) {
        std::cerr << "模型加載失敗: " << e.what() << std::endl;
        exit(-1);
    }
}

3.2 構建虛擬輸入張量

torch::Tensor create_dummy_input(int batch_size = 1) {
    // YOLOv5標準輸入尺寸為640x640
    return torch::ones({batch_size, 3, 640, 640}, 
                      torch::dtype(torch::kFloat32));
}

3.3 執行ONNX導出

#include <torch/csrc/jit/passes/onnx.h>

void export_to_onnx(torch::jit::script::Module& model, 
                   const torch::Tensor& dummy_input,
                   const std::string& output_path) {
    
    // 設置導出參數
    auto onnx_params = torch::onnx::ExportOptions();
    onnx_params.opset_version = 12;  // 使用ONNX opset 12
    
    // 執行導出
    torch::jit::ExportOutput export_output = torch::jit::export_onnx(
        model,
        {dummy_input},
        output_path,
        onnx_params,
        false,  // 不導出原始IR
        true    // 保持動態軸
    );
    
    std::cout << "ONNX模型已保存至: " << output_path << std::endl;
}

四、動態軸處理技巧

4.1 動態批次處理實現

void enable_dynamic_axis(torch::onnx::ExportOptions& options) {
    // 設置動態維度
    std::unordered_map<std::string, std::vector<int64_t>> dynamic_axes;
    dynamic_axes["input"] = {0};  // 批次維度動態
    dynamic_axes["output"] = {0}; // 輸出批次動態
    
    options.dynamic_axes = dynamic_axes;
}

4.2 多尺度輸入支持

void set_multiscale_support() {
    // 需要在模型結構中實現多尺度處理邏輯
    // 通常需要在Python端修改模型定義
    // C++端主要確保動態軸設置正確
}

五、后處理優化

5.1 輸出節點簡化

void simplify_output(torch::jit::script::Module& model) {
    // 通過圖優化簡化輸出節點
    torch::jit::pass::SimplifyOutput(model);
}

5.2 非極大抑制(NMS)集成

建議在導出前修改模型定義，將NMS作為模型的一部分：

# 在YOLOv5模型中添加NMS
model.model[-1].include_nms = True

六、完整C++實現示例

6.1 主程序實現

#include <iostream>
#include <torch/script.h>
#include <torch/csrc/jit/passes/onnx.h>

int main(int argc, char** argv) {
    if (argc != 3) {
        std::cerr << "用法: " << argv[0] << " <input.pt> <output.onnx>" << std::endl;
        return -1;
    }
    
    // 加載模型
    auto model = load_model(argv[1]);
    
    // 創建虛擬輸入
    auto dummy_input = create_dummy_input();
    
    // 配置導出選項
    torch::onnx::ExportOptions options;
    options.opset_version = 12;
    enable_dynamic_axis(options);
    
    // 執行導出
    export_to_onnx(model, dummy_input, argv[2]);
    
    std::cout << "轉換完成!" << std::endl;
    return 0;
}

6.2 CMake構建配置

cmake_minimum_required(VERSION 3.12)
project(yolov5_export)

set(CMAKE_CXX_STANDARD 14)

find_package(Torch REQUIRED)
find_package(ONNXRuntime REQUIRED)

add_executable(yolov5_export main.cpp)
target_link_libraries(yolov5_export ${TORCH_LIBRARIES})

七、常見問題解決方案

7.1 算子不支持錯誤

// 處理不支持的算子
void handle_unsupported_operators() {
    try {
        // 正常導出代碼
    } catch (const std::exception& e) {
        std::cerr << "錯誤: " << e.what() << std::endl;
        // 可能需要降低opset版本或自定義算子
    }
}

7.2 形狀推斷失敗

解決方案： 1. 檢查輸入張量形狀是否正確 2. 確保模型中沒有動態形狀操作 3. 嘗試固定批次大小

7.3 性能優化建議

// 啟用圖優化
void enable_optimizations(torch::jit::script::Module& model) {
    torch::jit::GraphOptimizerEnabledGuard guard(true);
    model = torch::jit::optimize_for_inference(model);
}

八、驗證轉換結果

8.1 ONNX Runtime驗證

#include <onnxruntime/core/session/onnxruntime_cxx_api.h>

void validate_onnx_model(const std::string& model_path) {
    Ort::Env env;
    Ort::SessionOptions session_options;
    Ort::Session session(env, model_path.c_str(), session_options);
    
    // 獲取輸入輸出信息
    auto input_info = session.GetInputTypeInfo(0);
    auto output_info = session.GetOutputTypeInfo(0);
    
    std::cout << "模型驗證通過!" << std::endl;
}

8.2 精度驗證方法

建議使用相同的測試數據分別運行原始PyTorch模型和轉換后的ONNX模型，對比輸出結果的差異。

九、進階話題

9.1 量化支持

void apply_quantization(const std::string& model_path) {
    // 需要ONNX Runtime的量化工具支持
    // 通常建議在Python端完成量化
}

9.2 多平臺部署考慮

1. 考慮目標平臺的算子支持情況
2. 針對不同硬件進行優化
3. 測試不同ONNX Runtime版本兼容性

結語

通過本文詳細介紹，我們了解了如何使用C++環境將YOLOv5模型轉換為ONNX格式。關鍵點包括： 1. 正確配置LibTorch和ONNX Runtime環境 2. 處理動態軸和特殊算子 3. 優化輸出結構和后處理 4. 驗證轉換結果的正確性

實際應用中，可能需要根據具體需求調整轉換參數。建議在轉換完成后進行充分的測試驗證，確保模型在目標部署環境中的性能和精度符合預期。

附錄

參考資源

1. YOLOv5官方倉庫
2. ONNX官方文檔
3. LibTorch C++ API文檔

工具推薦

1. Netron：可視化ONNX模型結構
2. ONNX-TensorRT：進一步轉換為TensorRT引擎
3. ONNX Simplifier：簡化復雜模型結構

”`