Интеграция On-Device ML модели (ONNX Runtime) для мобильного приложения

Интеграция On-Device ML модели (ONNX Runtime) для мобильного приложения

ONNX Runtime Mobile привлекателен одним аргументом: одна модель — обе платформы. Конвертировали PyTorch в ONNX, подключили onnxruntime-android и onnxruntime-objc, запускаете один и тот же .onnx файл. На практике разница в execution providers между iOS и Android всё равно требует платформо-специфичного кода, но сама модель едина.

Подготовка модели для мобиля

Стандартный ONNX экспорт из PyTorch:

import torch
import onnx
from onnxsim import simplify  # onnx-simplifier для оптимизации графа

model = MyModel(); model.eval()
dummy = torch.zeros(1, 3, 224, 224)

torch.onnx.export(
    model, dummy, "model.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

# Упрощение графа — убирает лишние reshape, transpose, делает граф чище
model_onnx = onnx.load("model.onnx")
model_simplified, check = simplify(model_onnx)
onnx.save(model_simplified, "model_simplified.onnx")

Для мобиля дополнительно — квантизация через onnxruntime.quantization:

from onnxruntime.quantization import quantize_dynamic, QuantType

quantize_dynamic(
    "model_simplified.onnx",
    "model_int8.onnx",
    weight_type=QuantType.QInt8
)
# Размер модели уменьшается в ~4× против FP32

Android: подключение и запуск

// build.gradle
implementation("com.microsoft.onnxruntime:onnxruntime-android:1.18.0")

// Создание сессии
val sessionOptions = OrtSession.SessionOptions().apply {
    // NNAPI Execution Provider для Android NPU/DSP
    addNnapi(NNAPIFlags.USE_FP16)  // FP16 режим в NNAPI
    // Или: addXnnpack(mapOf()) для XNNPACK (CPU SIMD)
    setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT)
    setIntraOpNumThreads(4)
}

val env = OrtEnvironment.getEnvironment()
val session = env.createSession(
    context.assets.open("model_simplified.onnx").readBytes(),
    sessionOptions
)

// Инференс
val inputTensor = OnnxTensor.createTensor(
    env,
    FloatBuffer.wrap(preprocessedArray),
    longArrayOf(1, 3, 224, 224)
)

val results = session.run(mapOf("input" to inputTensor))
val outputArray = (results["output"]?.value as Array<FloatArray>)[0]

// Освобождение ресурсов — обязательно
inputTensor.close()
results.close()

Утечки через незакрытые OnnxTensor и OrtSession.Result — частая проблема. В Kotlin используем use {} блок: results.use { ... }.

iOS: ObjC/Swift интеграция

// Package.swift или Podfile: pod 'onnxruntime-objc'
import onnxruntime_objc

// Настройка
let env = try ORTEnv(loggingLevel: ORTLoggingLevel.warning)
let options = try ORTSessionOptions()
try options.setIntraOpNumThreads(4)
// На iOS — CoreML Execution Provider
try options.appendCoreMLExecutionProvider(withFlags: [.enableOnSubgraphs])

let session = try ORTSession(
    env: env,
    modelPath: Bundle.main.path(forResource: "model_simplified", ofType: "onnx")!,
    sessionOptions: options
)

// Подготовка входа
let inputShape: [NSNumber] = [1, 3, 224, 224]
let inputData = Data(bytes: preprocessedFloats, count: preprocessedFloats.count * MemoryLayout<Float>.size)
let inputTensor = try ORTValue(
    tensorData: NSMutableData(data: inputData),
    elementType: .float,
    shape: inputShape
)

let outputs = try session.run(
    withInputs: ["input": inputTensor],
    outputNames: ["output"],
    runOptions: nil
)

let outputTensor = outputs["output"]!
let outputData = try outputTensor.tensorData() as Data
let floats = outputData.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) }

appendCoreMLExecutionProvider на iOS 13+ делегирует поддерживаемые операции в Core ML, что даёт доступ к ANE. Операции, которые Core ML не поддерживает, автоматически выполняются на CPU. Это не так быстро, как нативный Core ML с полным графом на ANE, но удобно для быстрого кросс-платформенного деплоя.

Когда ONNX Runtime лучше нативных форматов

Да: прототипирование, кросс-платформенная модель, модели с нестандартными операциями которые coremltools не конвертирует, частые обновления модели без пересборки конвертационного пайплайна.

Нет: максимальная производительность на одной платформе. Нативный Core ML на iOS с полным ANE-ускорением обычно быстрее ORT+CoreML EP на 20–40%. TFLite + GPU Delegate на Android в ряде случаев быстрее ORT+NNAPI. Если модель разворачивается только на одной платформе и производительность критична — берите нативный формат.

Отладка несовместимых операций

# Проверить, какие операции поддерживает NNAPI Execution Provider
python -m onnxruntime.tools.check_nnapi_supported_ops --model model.onnx

# Если операция не поддерживается — она выполнится на CPU (fallback)
# Это не краш, но может обнулить всё ускорение от NNAPI

Для идентификации узких мест: ORT Profiling API записывает время каждого оператора. Включается через options.enableProfiling("ort_profile") — генерирует JSON, открываемый в Chrome chrome://tracing.

Процесс

Экспорт и упрощение ONNX-графа → квантизация → интеграция на iOS и Android с подходящими EP → профилирование и сравнение с нативными форматами → финальный выбор runtime под продакшн.

Ориентиры по срокам

Базовая кросс-платформенная интеграция ONNX Runtime — 2–3 недели. С оптимизацией EP, профилированием, тестированием на парке устройств — 4–6 недель.