Android 手机 ONNX 推理加速实战：从模型优化到性能调优

1次阅读

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

在移动端部署 ONNX 模型时，开发者常面临推理延迟高、内存占用大等问题。本文将针对 Android 平台，系统性地介绍 ONNX 模型量化、图优化、多线程推理等加速技术，并提供可落地的代码实现。通过本文，开发者将掌握一套完整的移动端模型优化方案，显著提升推理性能。

在 Android 设备上运行 ONNX 模型时，通常会遇到以下几个主要瓶颈：

CPU 算力有限 ：移动端 CPU 性能较弱，尤其是浮点运算能力
内存压力 ：大模型加载导致内存占用高，容易触发 OOM
线程调度 ：默认单线程推理无法充分利用多核 CPU
功耗限制 ：长时间高负载运行会导致降频

Android 平台主要有以下几种推理引擎选择：

ONNX Runtime：官方支持，功能全面，优化较好
NNAPI：系统级 API，能利用硬件加速
TFLite：轻量级，但对 ONNX 模型需要转换

对于 ONNX 模型，我们推荐使用 ONNX Runtime，原因如下：

原生支持 ONNX 格式
提供多种优化选项
跨平台一致性更好

量化是减少模型大小和提高推理速度的最有效方法之一。ONNX 支持两种量化方式：

静态量化 ：
需要校准数据集
离线完成量化
推理时完全使用整型运算
动态量化 ：
无需校准数据
运行时自动量化
部分运算仍为浮点

量化实现步骤：

# 静态量化示例代码
from onnxruntime.quantization import quantize_static, CalibrationDataReader

# 1. 准备校准数据
class MyCalibrationDataReader(CalibrationDataReader):
    def __init__(self):
        # 实现数据读取逻辑
        pass

# 2. 执行量化
quantize_static(
    'float_model.onnx',
    'quant_model.onnx',
    MyCalibrationDataReader())

ONNX Runtime 提供了多种图优化技术：

节点融合：将多个算子合并为一个
常量折叠：预先计算静态值
冗余节点消除：删除无用节点

启用方法：

// 在 Android 中初始化 ONNX Runtime 时启用优化
OrtSession.SessionOptions options = new OrtSession.SessionOptions();
options.setOptimizationLevel(OptimizationLevel.ALL_OPT);

合理配置线程池可以显著提升性能：

根据 CPU 核心数设置线程数（通常为 CPU 核心数 -1）
避免线程过多导致上下文切换开销
注意线程安全，特别是模型共享时

完整示例代码：

// native-lib.cpp
#include <jni.h>
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>

// 初始化 ONNX Runtime
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntime");
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(4); // 设置线程数

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

// 准备输入输出
std::vector<const char*> input_names = {"input"};
std::vector<const char*> output_names = {"output"};

// 推理函数
extern "C" JNIEXPORT jfloatArray JNICALL
Java_com_example_onnxdemo_MainActivity_runInference(
    JNIEnv* env,
    jobject /* this */,
    jfloatArray input) {

    // 获取输入数据
    jfloat* input_data = env->GetFloatArrayElements(input, nullptr);

    // 创建输入 Tensor
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

    std::vector<int64_t> input_shape = {1, 3, 224, 224}; // 示例 shape
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_data, 3*224*224, input_shape.data(), input_shape.size());

    // 执行推理
    auto output_tensors = session.Run(Ort::RunOptions{nullptr},
        input_names.data(), &input_tensor, 1,
        output_names.data(), 1);

    // 处理输出
    float* output_data = output_tensors[0].GetTensorMutableData<float>();
    // ...
}

优化前后对比数据示例：