Rooftop

ONNX Runtime

7分钟阅读

Onnx是一种模型交换格式，可以理解为“训练框架和推理框架之间的通用中间格式”：

  • 训练通常是PyTorch、TensorFlow里做。
  • 导出时转成.onnx。
  • 线上推理是由ONNX Runtime去加载和执行。
  • 它解决的是训练框架和部署环境解耦的问题。训练和部署不必绑定在同一个框架里。

ONNX、ONNX Runtime、App代码三者关系

  • ONNX：模型交换格式，用来描述模型结构、权重、输入输出。
  • ONNX Runtime：模型执行引擎，负责加载.onnx并完成推理。
  • App代码：负责业务数据的预处理、调用推理、结果后处理。

可以简单理解为： 业务数据 → App预处理 → ONNX Runtime推理 → App后处理 → 业务结果

什么是张量（Tensor）？ 张量可以先理解为“带shape的多维数组”。 例如人脸模型输入常见是[1, 3, 112, 112]：

  • 1：batch size，一次输入1张图。
  • 3：3个颜色通道。
  • 112,112：图片宽高。

除了shape，还要关注数据类型dtype，例如常见的float32。

ONNX Runtime核心对象

  1. OrtEnvironment 运行时的全局环境，通常整个进程里复用，不要频繁创建销毁。
  2. OrtSession 理解为“已经加载好的模型实例”。真实推理基本都通过它来做。
  3. OnnxTensor 它是传给模型的输入张量。这里底层数据是FloatBuffer.wrap(inputData)。
  4. OrtSession.Result 它持有输出张量，使用后要close()。

什么是预处理 / 后处理？

  • 预处理：把业务数据转换成模型能吃的输入格式。
  • 推理：模型根据输入做计算。
  • 后处理：把模型输出转换成业务可用结果。

在图片模型里，预处理通常包括：

  • resize
  • RGB/BGR通道转换
  • NCHW/NHWC布局转换
  • 像素归一化

在这个人脸比对例子里，后处理通常包括：

  • 对embedding做L2归一化
  • 计算余弦相似度
  • 按阈值判断是否相似

ONNX在端上（Android）的典型流程：

  1. 加载模型
    • 从assets读取models/mobilefacenet.onnx（这里以人脸识别模型MobileFaceNet为例）。
    • 创建OrtEnvironment。
    • 配置OrtSession.SessionOptions。
    • 创建OrtSession。
  2. 预处理输入图片
    • 把Bitmap缩放到112x112。
    • 转成float[]。
    • 按模型要求组织成张量shape：[1, 3, 112, 112]。
  3. 执行推理
    • 取输入名。
    • 构造Map<String, OnnxTensor>。
    • ortSession.run(inputs)。
  4. 读取输出
    • 输出是float[][]。
    • 取第一个batch的向量output[0]。
    • 再做一次L2 normalize。
  5. 做相似度比较
    • 本质是余弦相似度。
    • 因为前面做了L2归一化，所以dot(a,b)就基本等于cosine similarity。
    • 把[-1,1]映射到[0,1]。

为什么输入是[1, 3, 112, 112]？ 这个和模型定义有关：

  • 1：batch size。
  • 3：RGB三通道。
  • 112,112：MobileFaceNet常见输入尺寸。

除了shape本身，还要确认两件事：

  • 数据布局是不是NCHW，不要和NHWC搞混。
  • 通道顺序是不是RGB，不要误用成BGR。

为什么预处理最容易出问题？ 模型能跑通，不代表结果正确。 实际接入里，最重要的不是把模型跑起来，而是保证“输入预处理 + 输出后处理”与模型训练时保持一致。

最常见的错误包括：

  • 输入尺寸不对。
  • 通道顺序错了（RGB/BGR）。
  • 张量布局错了（NCHW/NHWC）。
  • 归一化方式错了，比如[0,1]、[-1,1]、mean/std不一致。

怎么接入GPU？ 尝试options.addNnapi()。 这里并不是“直接调用GPU API”，而是让ONNX Runtime把模型交给Android NNAPI，底层可能落到DSP/NPU/GPU/CPU，取决于设备驱动和算子支持情况。

SessionOptions本质上是在配置推理策略，常见包括：

  • 图优化等级。
  • 线程数。
  • 执行后端（CPU / NNAPI）。

这些配置通常影响性能、功耗、兼容性，但不改变模型本身的语义。

为什么要warmup？ 首次推理往往比后续慢，常见原因包括：

  • 模型加载。
  • 内存分配。
  • 图优化。
  • NNAPI编译或初始化。

所以性能测试时，通常要区分：

  • 首次推理耗时。
  • 稳定运行后的平均耗时。

资源释放为什么重要？

  • OrtSession.close()
  • OrtEnvironment.close()

在Android上做端侧推理，内存和native资源都比较敏感。 OnnxTensor、Result、Session这些对象很多都带native句柄，不及时释放容易出现：

  • 内存上涨。
  • 句柄泄漏。
  • 长时间运行后崩溃。

ONNX Runtime常见坑

  1. 模型能跑，但结果不对。 常见原因不是Runtime本身，而是预处理不一致，比如尺寸、通道顺序、归一化方式错了。
  2. shape不匹配。 创建OnnxTensor时，数据长度和shape必须对应，否则运行时会报错。
  3. 输入输出名写错。 多输入、多输出模型里，不能想当然只取第一个输入或输出。
  4. 首次推理很慢。 可能是模型初始化、图优化、NNAPI编译带来的开销，通常需要warmup后再看稳定耗时。
  5. native资源没释放。 OnnxTensor、Result、Session等对象可能持有native资源，不及时close()会导致内存问题。