发布于 由admin

C++ ONNX Runtime / LibTorch C++ API：高性能 AI 模型部署与推理优化

哈喽，各位好！

今天咱们来聊聊怎么让 AI 模型跑得飞快，尤其是在 C++ 环境下。咱们的主题是“C++ ONNX Runtime / LibTorch C++ API：高性能 AI 模型部署与推理优化”。 这可不是纸上谈兵，咱们会撸起袖子，直接上代码，保证让大家看得明白，学得会。

一、模型部署与推理的必要性：为啥要折腾 C++？

你可能觉得，Python 写起来多爽啊，为啥还要费劲巴拉地用 C++？原因很简单：速度！

  • 性能至上： C++ 编译后直接生成机器码，执行效率比解释型语言 Python 高得多。在对延迟要求高的场景，比如实时语音识别、自动驾驶，C++ 简直是救命稻草。
  • 资源限制： 嵌入式设备、移动设备等资源有限，C++ 可以更精细地控制内存和 CPU 使用，让模型在“蜗居”里也能跑起来。
  • 现有系统集成： 很多传统系统都是 C++ 写的，直接用 C++ 部署 AI 模型，可以无缝集成，避免不必要的麻烦。

二、ONNX Runtime：模型跨平台运行的利器

ONNX (Open Neural Network Exchange) 是一种开放的模型格式，旨在让不同的 AI 框架（PyTorch, TensorFlow, MXNet 等）能够共享模型。 ONNX Runtime 是微软开源的一个跨平台推理引擎，可以加载 ONNX 模型，并在各种硬件平台上高效运行。

  • 跨平台： Windows, Linux, macOS, Android, iOS… ONNX Runtime 几乎无所不能。
  • 硬件加速： 支持 CPU, GPU (CUDA, TensorRT), 以及各种专用加速器。
  • 优化： 内置了很多优化技术，比如算子融合、常量折叠等，让模型跑得更快。

2.1 ONNX Runtime 的基本使用

先来个简单的例子，让大家感受一下 ONNX Runtime 的威力。

  1. 安装 ONNX Runtime:

    # CPU 版本
pip install onnxruntime
# GPU 版本 (需要先安装 CUDA)
pip install onnxruntime-gpu

  2. Python 模型导出为 ONNX 格式:

    import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()
dummy_input = torch.randn(1, 10)  # 1 个 batch, 10 个特征
torch.onnx.export(model, dummy_input, "simple_model.onnx", verbose=True,
                    input_names=['input'], output_names=['output'])

    这段代码定义了一个简单的线性模型，然后用 torch.onnx.export 函数把它导出为 ONNX 格式。dummy_input 是一个假的输入，用来告诉 ONNX Runtime 模型的输入形状。 input_names 和 output_names分别指定了输入和输出的名称，后面 C++ 代码会用到。

  3. C++ 代码加载 ONNX 模型并推理:

    #include <iostream>
#include <vector>
#include <onnxruntime_cxx_api.h>

int main() {
    // 1. 创建 ONNX Runtime 环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "my-onnx-runtime");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(4); // 设置线程数

    // 2. 加载 ONNX 模型
    Ort::Session session(env, "simple_model.onnx", session_options);

    // 3. 获取输入和输出信息
    Ort::AllocatorWithDefaultOptions allocator;
    const char* input_name = session.GetInputName(0, allocator);
    const char* output_name = session.GetOutputName(0, allocator);
    std::cout << "Input Name: " << input_name << std::endl;
    std::cout << "Output Name: " << output_name << std::endl;

    // 4. 创建输入数据
    std::vector<float> input_data(10, 0.5f); // 10 个特征，都初始化为 0.5
    std::vector<int64_t> input_shape = {1, 10}; // batch_size = 1, feature_size = 10

    // 5. 创建 ONNX Runtime 输入张量
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        input_data.data(), input_data.size(), input_shape.data(), input_shape.size());
    std::cout << "input_tensor created" << std::endl;
    // 6. 推理
    std::vector<Ort::Value> output_tensors = session.Run(
        {nullptr, &input_name, &input_tensor, 1}, // 输入
        {nullptr, &output_name, 1}  // 输出
    );

    // 7. 获取输出数据
    float* output_data = output_tensors[0].GetTensorMutableData<float>();
    auto output_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();

    std::cout << "Output Shape: ";
    for(auto dim : output_shape){
        std::cout << dim << " ";
    }
    std::cout << std::endl;

    std::cout << "Output Data: ";
    for (int i = 0; i < output_shape[1]; ++i) {
        std::cout << output_data[i] << " ";
    }
    std::cout << std::endl;

    return 0;
}

    这段 C++ 代码做了这些事情：

    • 创建 ONNX Runtime 环境和 Session: 相当于初始化 ONNX Runtime 引擎，并加载 ONNX 模型。
    • 获取输入/输出名称: ONNX Runtime 需要知道输入和输出张量的名字，才能正确地进行推理。
    • 创建输入张量: 把 C++ 的 std::vector 数据转换为 ONNX Runtime 的 Ort::Value 类型。
    • 推理: 调用 session.Run 函数，传入输入张量，得到输出张量。
    • 获取输出数据: 把 ONNX Runtime 的输出张量转换为 C++ 的 float* 指针，就可以访问输出数据了。

    编译运行:

    g++ -o onnx_example onnx_example.cpp -I/path/to/onnxruntime/include -L/path/to/onnxruntime/lib -lonnxruntime
./onnx_example

    记得把 /path/to/onnxruntime 替换成你实际的 ONNX Runtime 安装路径。

2.2 ONNX Runtime 优化技巧

光能跑起来还不够，咱们的目标是跑得飞快！ONNX Runtime 提供了很多优化选项，可以根据你的硬件和模型进行调整。

  • 线程数: session_options.SetIntraOpNumThreads(4) 设置了线程数，可以充分利用多核 CPU。

  • 优化级别: session_options.SetGraphOptimizationLevel(ORT_ENABLE_EXTENDED) 可以开启更高级别的图优化。

  • CUDA 加速: 如果你有 NVIDIA GPU，可以开启 CUDA 加速，让模型跑得更快。

    // CUDA 加速
#ifdef USE_CUDA
OrtCUDAProviderOptions cuda_options;
session_options.AppendExecutionProvider_CUDA(cuda_options);
#endif

  • TensorRT 加速: TensorRT 是 NVIDIA 的高性能推理引擎，可以进一步优化 ONNX 模型。

    // TensorRT 加速
#ifdef USE_TENSORRT
OrtTensorRTProviderOptions tensorrt_options;
session_options.AppendExecutionProvider_TensorRT(tensorrt_options);
#endif

    表格：ONNX Runtime 优化选项

    选项 描述
    SetIntraOpNumThreads 设置线程数，充分利用多核 CPU。
    SetGraphOptimizationLevel 设置图优化级别，可以开启更高级别的图优化，比如算子融合、常量折叠等。
    AppendExecutionProvider_CUDA 开启 CUDA 加速，利用 NVIDIA GPU 进行推理。 需要先安装 CUDA toolkit，并且编译 ONNX Runtime 的 CUDA 版本。
    AppendExecutionProvider_TensorRT 开启 TensorRT 加速，利用 NVIDIA TensorRT 引擎进行推理。 需要先安装 TensorRT，并且编译 ONNX Runtime 的 TensorRT 版本。 TensorRT 会对模型进行更深度的优化，通常可以获得更好的性能。
    SetSessionExecutionMode 设置会话执行模式，可以是顺序执行或者并行执行。 并行执行可以充分利用多核 CPU，提高推理速度。

三、LibTorch C++ API：PyTorch 模型原生部署

如果你用的是 PyTorch，那么 LibTorch C++ API 可能是更自然的选择。 LibTorch 是 PyTorch 的 C++ 版本，可以直接加载 PyTorch 模型，并在 C++ 环境下进行推理。

  • 原生支持： 直接加载 PyTorch 模型，不需要转换格式。
  • 灵活性： 可以自定义 C++ 算子，扩展 PyTorch 的功能。
  • 调试方便： 可以用 C++ 调试器调试 PyTorch 模型。

3.1 LibTorch 的基本使用

  1. 安装 LibTorch:

    从 PyTorch 官网下载 LibTorch 的 C++ 包，解压到指定目录。

  2. Python 模型导出为 TorchScript 格式:

    import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()
model.eval() # 非常重要！ 必须设置为 eval 模式
dummy_input = torch.randn(1, 10)

# 方法一：Tracing based
traced_script_module = torch.jit.trace(model, dummy_input)
traced_script_module.save("simple_model.pt")

# 方法二： Scripting based
# scripted_module = torch.jit.script(model)
# scripted_module.save("simple_model.pt")

    这段代码和 ONNX 的例子很像，也是定义了一个简单的线性模型，然后用 torch.jit.trace 函数把它导出为 TorchScript 格式。 model.eval() 非常重要， 必须设置为 eval 模式，告诉 PyTorch 模型处于推理模式，关闭 dropout 等训练相关的操作。
    TorchScript 有两种方式导出模型: tracing 和 scripting. Tracing 适合于结构简单的模型，Scripting 适合于包含复杂控制流的模型。

  3. C++ 代码加载 TorchScript 模型并推理:

    #include <iostream>
#include <torch/script.h> // One-stop header.
#include <torch/torch.h>

int main() {
    // 1. 加载 TorchScript 模型
    torch::jit::Module module = torch::jit::load("simple_model.pt");

    // 2. 创建输入张量
    std::vector<float> input_data(10, 0.5f);
    torch::Tensor input_tensor = torch::from_blob(input_data.data(), {1, 10}, torch::kFloat32);

    // 3. 推理
    torch::Tensor output_tensor = module.forward({input_tensor}).toTensor();

    // 4. 获取输出数据
    float* output_data = output_tensor.data_ptr<float>();
    auto output_shape = output_tensor.sizes();

    std::cout << "Output Shape: ";
    for(auto dim : output_shape){
        std::cout << dim << " ";
    }
    std::cout << std::endl;

    std::cout << "Output Data: ";
    for (int i = 0; i < output_shape[1]; ++i) {
        std::cout << output_data[i] << " ";
    }
    std::cout << std::endl;

    return 0;
}

    这段 C++ 代码做了这些事情：

    • 加载 TorchScript 模型: torch::jit::load 函数加载 TorchScript 模型。
    • 创建输入张量: 把 C++ 的 std::vector 数据转换为 PyTorch 的 torch::Tensor 类型。
    • 推理: 调用 module.forward 函数，传入输入张量，得到输出张量。
    • 获取输出数据: 把 PyTorch 的输出张量转换为 C++ 的 float* 指针，就可以访问输出数据了。

    编译运行:

    g++ -o libtorch_example libtorch_example.cpp -I/path/to/libtorch/include -I/path/to/libtorch/include/torch/csrc/api/include -L/path/to/libtorch/lib -ltorch -ltorch_cpu -lc10
./libtorch_example

    记得把 /path/to/libtorch 替换成你实际的 LibTorch 安装路径。

3.2 LibTorch 优化技巧

LibTorch 也提供了一些优化选项，可以提高推理速度。

  • 设置线程数:

    torch::set_num_threads(4);

  • CUDA 加速: 如果你有 NVIDIA GPU，可以把模型和输入数据都放到 GPU 上进行推理。

    torch::Device device(torch::kCUDA);
module.to(device);
input_tensor = input_tensor.to(device);
torch::Tensor output_tensor = module.forward({input_tensor}).toTensor();

  • Just-In-Time (JIT) 编译: LibTorch 的 JIT 编译器可以对模型进行优化，提高推理速度。

    torch::jit::Module module = torch::jit::load("simple_model.pt");
module.optimize_for_inference(); // 开启优化

    表格：LibTorch 优化选项

    选项 描述
    torch::set_num_threads 设置线程数，充分利用多核 CPU。
    torch::Device 指定设备类型，可以是 CPU (torch::kCPU) 或者 GPU (torch::kCUDA)。 如果有 NVIDIA GPU，可以把模型和输入数据都放到 GPU 上进行推理，提高推理速度。
    module.optimize_for_inference() 开启 JIT 编译器的优化，可以对模型进行优化，比如算子融合、常量折叠等，提高推理速度。
    torch::NoGradGuard 在推理过程中，不需要计算梯度，可以关闭梯度计算，减少内存占用，提高推理速度。 使用 torch::NoGradGuard no_grad; 可以创建一个 no_grad 上下文，在这个上下文中，所有的操作都不会计算梯度。

四、性能对比与选择建议

ONNX Runtime 和 LibTorch C++ API 各有优缺点，选择哪个取决于你的具体需求。

  • ONNX Runtime:

    • 优点: 跨平台性好，支持多种硬件加速器，优化选项丰富。
    • 缺点: 需要把模型转换为 ONNX 格式，可能会引入一些精度损失。

  • LibTorch C++ API:

    • 优点: 原生支持 PyTorch 模型，使用方便，调试方便。
    • 缺点: 跨平台性不如 ONNX Runtime，对硬件加速器的支持不如 ONNX Runtime 丰富。

    表格：ONNX Runtime vs LibTorch C++ API

    特性 ONNX Runtime LibTorch C++ API
    模型格式 ONNX TorchScript
    跨平台性 好，支持多种操作系统和硬件平台 较好，主要支持 Linux, Windows, macOS
    硬件加速 支持 CPU, GPU (CUDA, TensorRT), 专用加速器 支持 CPU, GPU (CUDA)
    易用性 需要转换模型格式，配置相对复杂 直接加载 PyTorch 模型，使用方便
    调试 相对困难 方便，可以使用 C++ 调试器调试 PyTorch 模型
    适用场景 需要跨平台部署，或者需要利用特定硬件加速器的场景 使用 PyTorch，并且需要在 C++ 环境下进行高性能推理的场景

选择建议:

  • 如果你的模型需要在多种平台上部署，并且需要利用特定硬件加速器，那么 ONNX Runtime 是更好的选择。
  • 如果你的模型是 PyTorch 写的，并且只需要在少数平台上部署，那么 LibTorch C++ API 可能是更方便的选择。

五、总结与展望

今天咱们聊了 C++ ONNX Runtime 和 LibTorch C++ API 这两种高性能 AI 模型部署方案。 掌握这些技术，可以让你在 C++ 环境下也能轻松驾驭 AI 模型，让它们跑得飞快！

当然，AI 模型部署和推理优化是一个持续发展的领域。 随着硬件和算法的不断进步，未来会有更多更高效的部署方案涌现。 让我们一起学习，共同进步！

希望今天的分享对大家有所帮助！ 谢谢大家！

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注