使用C++在端上运行机器学习模型

上个季度加入了公司的 AI 团队，负责的工作除了常规的服务端开发之外，还包括了在端上运行机器学习模型。考虑到性能因素，没有使用 Python 或 Node.js 等运行时，而是直接使用 C++来执行模型推理逻辑。在这里分享一些经验和教训。

模型格式

我们使用的模型格式是 onnx。onnx 是一种开放的模型格式，使用一个二进制文件存储模型的结构和参数，可以在不同的平台上使用不同的运行时加载和执行。目前主流的深度学习框架都支持导出 onnx 模型，例如 PyTorch、TensorFlow 等。

模型运行时

对于常见的网络结构的模型，可以使用 OpenCV 的 dnn 模块来加载和执行。OpenCV 是一个开源的计算机视觉库，提供了很多计算机视觉相关的功能，其中 dnn 模块提供了加载和执行深度学习模型的功能。

对于一些使用了自定义算子的模型，OpenCV 的 dnn 模块可能无法加载，这时候可以使用 onnxruntime 来加载和执行模型。onnxruntime 是微软开源的一个 onnx 运行时，支持加载和执行 onnx 模型，支持 CPU 和 GPU 加速。小小的吐槽一下，与 Python 的 onnxruntime sdk 相比，C++ 的 onnxruntime sdk 的接口有够难用。

Windows 上使用 onnxruntime 的一个坑

Windows 10 和 Windows 11 在系统库里内置了一个较旧版本的 onnx runtime，如果使用了较新的 onnxruntime，可能会出现加载模型失败的问题。一个解决方式是在代码中指定加载动态链接库的路径，并通过 onnx runtime 的 api 指定加载的版本。

#include "Windows.h"
#define ORT_API_MANUAL_INIT // 定义这个宏可以禁止 onnxruntime 自动加载
#include <onnxruntime_cxx_api.h>

int main(){
  auto mod = LoadLibrary(L"onnxruntime1.16.1.dll");
  auto OrtGetApiBase =
      (const OrtApiBase *(*)(void))GetProcAddress(mod, "OrtGetApiBase");
  Ort::InitApi(OrtGetApiBase()->GetApi(ORT_API_VERSION));

  // ...
}

模型推理的一般流程

加载模型

OpenCV 的 dnn 模块使用 cv::dnn::readNetFromONNX 函数加载模型
onnxruntime 使用 Ort::SessionOptions 和 Ort::Session 类加载模型

预处理输入数据

对于两种运行时，都可以使用 OpenCV 的 dnn 模块中 cv::dnn::blobFromImage 函数将输入图像转换为模型需要的格式，可以在一个函数中完成矩阵转换、交换通道、归一化等操作。

对于 onnxruntime，可以将处理好的 cv::Mat 转换为 onnxruntime 的 Tensor 格式。

auto input_tensor = cv::dnn::blobFromImage(
                input, 1.0 / 255, cv::Size(320, 320), cv::Scalar(), true);
std::vector<float> input_values;
input_values.assign((float *)input_tensor.ptr(),
                    (float *)input_tensor.ptr() +
                        1 * 3 * 320 * 320);

auto memoryInfo = Ort::MemoryInfo::CreateCpu(
				OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeCPU);
std::vector<Ort::Value> input_tensors;
input_tensors.push_back(Ort::Value::CreateTensor<float>(memoryInfo, input_values.data(), input_tensor_size,         input_node_dims.data(), input_node_dims.size()));

执行推理

OpenCV 的 dnn 模块使用 cv::dnn::Net::forward 函数执行推理
onnxruntime 使用 Ort::Session::Run 函数执行推理

后处理输出数据

对于两种运行时，都可以使用 OpenCV 的 dnn 模块中 cv::dnn::NMSBoxes 函数完成非极大值抑制等操作。
针对不同的任务类型（例如分类模型和回归模型），后处理的逻辑也会略有不同。

模型推理的性能优化

最直接的方式是使用 GPU 加速，对于 onnxruntime，可以通过 Ort::SessionOptions::AppendExecutionProvider_CUDA 函数添加 CUDA 加速。
实际上 Windows 平台也支持其他的运行时，例如 DirectML。

定义一个通用的模型推理接口

可以使用 C++的模板，std::function 等特性，定义一个通用的模型推理接口，从而避免为每一个模型都定义一大堆的类和函数。

#pragma once
#include <functional>
#include <onnxruntime_cxx_api.h>

namespace umr {
template <typename In, typename Out> struct ModelConfig {
  std::wstring model_path;
  std::function<std::vector<float>(In)> input_transform;
  std::function<Out(std::vector<float>, In)> output_transform;
};

template <typename In, typename Out> class UniversalModelRunner {
private:
  std::shared_ptr<Ort::Env> env;
  std::shared_ptr<Ort::Session> session;
  std::function<std::vector<float>(In)> input_transform;
  std::function<Out(std::vector<float>, In)> output_transform;

public:
  UniversalModelRunner(ModelConfig<In, Out> &&config)；
  Out run(In input);
}

上述代码中，ModelConfig 定义了模型的路径，输入数据的预处理函数，输出数据的后处理函数。UniversalModelRunner 是一个通用的模型推理类，使用模板参数 In 和 Out 来定义输入和输出的数据类型，使用 ModelConfig 来初始化模型。

在实际使用中，直接根据定义好的模型配置来初始化模型即可。

auto runner = UniversalModelRunner<cv::Mat, std::pair<int, float>>(
    ModelConfig<cv::Mat, std::vector<Box>>{
        L"D:/models/cls.onnx",
        [](cv::Mat input) {
            auto input_tensor = cv::dnn::blobFromImage(
                input, 1.0 / 255, cv::Size(320, 320), cv::Scalar(), true);
            std::vector<float> input_values;
            input_values.assign((float *)input_tensor.ptr(),
                                (float *)input_tensor.ptr() +
                                    1 * 3 * 320 * 320);
            return input_values;
        },
        [](std::vector<float> output_values, cv::Mat input) {
            auto maxElement = std::max_element(output_values.begin(),
                                                output_values.end());
            auto maxIndex = std::distance(output_values.begin(), maxElement);
            return std::make_pair(maxIndex, *maxElement);
        },
});

auto result = runner.run(input);
std::cout << result.first << " " << result.second << std::endl;

上述代码中，我们定义了一个模型推理类，用于执行分类模型的推理，输入数据是一个 cv::Mat 类型的图像，输出数据是一个 std::pair<int, float> 类型的结果，其中 int 表示分类的类别，float 表示分类的概率。

总结

尽管大语言模型是近期 AI 浪潮的主角，但受限于成本、性能与隐私安全等因素，端上的小语言模型也是一个重要的研究方向。在这篇文章中，我们介绍了如何使用 C++ 在端上运行机器学习模型，以及如何定义一个通用的模型推理接口，希望能够对读者有所帮助。