Интеграция On-Device ML модели (TensorFlow Lite) для офлайн AI в Android-приложении

Интеграция On-Device ML модели (TensorFlow Lite) для офлайн AI в Android-приложении

TensorFlow Lite — стандарт де-факто для запуска ML-моделей на Android. Но «добавить tflite файл в assets» — это только начало. Реальная интеграция включает выбор делегата ускорения, управление памятью буферов, обработку несовместимостей между устройствами и тестирование числовой точности.

Конвертация модели в TFLite

Из PyTorch через ONNX:

# PyTorch → ONNX
python -c "
import torch, onnx
model = MyModel(); model.eval()
torch.onnx.export(model, torch.zeros(1,3,224,224), 'model.onnx',
    opset_version=17, input_names=['input'], output_names=['output'])
"

# ONNX → TFLite через onnx-tf
pip install onnx-tf tensorflow
onnx-tf convert -i model.onnx -o model_tf
tflite_convert --saved_model_dir=model_tf --output_file=model.tflite

Или напрямую из TensorFlow SavedModel:

converter = tf.lite.TFLiteConverter.from_saved_model("model_tf")
converter.optimizations = [tf.lite.Optimize.DEFAULT]  # динамическая квантизация FP16
converter.target_spec.supported_types = [tf.float16]  # для GPU delegate
tflite_model = converter.convert()
with open("model_fp16.tflite", "wb") as f:
    f.write(tflite_model)

Делегаты ускорения: что выбрать

// GPU Delegate — наиболее универсальный
import org.tensorflow.lite.gpu.GpuDelegate
import org.tensorflow.lite.gpu.CompatibilityList

val compatList = CompatibilityList()
val options = Interpreter.Options()

if (compatList.isDelegateSupportedOnThisDevice) {
    val delegateOptions = compatList.bestOptionsForThisDevice
    options.addDelegate(GpuDelegate(delegateOptions))
} else {
    // Fallback: NNAPI или CPU с XNNPACK
    options.setUseNNAPI(true)
    options.setUseXNNPACK(true)
}
options.setNumThreads(4)

val interpreter = Interpreter(
    FileUtil.loadMappedFile(context, "model_fp16.tflite"),
    options
)

NNAPI на практике нестабилен: на одних устройствах даёт 5× ускорение, на других — краш с NNAPIDelegate: Failed to invoke the model из-за несовместимых операций. Обязательно — try/catch с fallback на CPU:

try {
    options.setUseNNAPI(true)
    interpreter = Interpreter(modelBuffer, options)
    // Тестовый прогон для проверки
    interpreter.run(testInput, testOutput)
} catch (e: Exception) {
    Log.w("ML", "NNAPI failed, falling back to CPU: ${e.message}")
    options.setUseNNAPI(false)
    interpreter = Interpreter(modelBuffer, options)
}

Управление буферами: ByteBuffer vs TensorBuffer

Прямое управление ByteBuffer — быстрее, но многословно. TensorBuffer из org.tensorflow.lite.support — удобнее:

// Через TFLite Support Library (рекомендую)
val imageProcessor = ImageProcessor.Builder()
    .add(ResizeOp(224, 224, ResizeOp.ResizeMethod.BILINEAR))
    .add(NormalizeOp(127.5f, 127.5f))  // нормализация [-1, 1]
    .build()

val tensorImage = TensorImage(DataType.FLOAT32)
tensorImage.load(bitmap)
val processedImage = imageProcessor.process(tensorImage)

// Запуск
val outputBuffer = TensorBuffer.createFixedSize(intArrayOf(1, 1000), DataType.FLOAT32)
interpreter.run(processedImage.buffer, outputBuffer.buffer)

// Результат
val probabilities = outputBuffer.floatArray
val topIndex = probabilities.indices.maxByOrNull { probabilities[it] } ?: -1

ResizeOp на CPU — неожиданно медленный для больших изображений (Full HD → 224×224 занимает 20–40 мс). Альтернатива: предварительный ресайз через Bitmap.createScaledBitmap() или через RenderScript (deprecated) / Camera2 output size.

CameraX интеграция

val imageAnalyzer = ImageAnalysis.Builder()
    .setTargetResolution(Size(640, 480))
    .setBackpressureStrategy(ImageAnalysis.STRATEGY_KEEP_ONLY_LATEST)  // не копим очередь
    .build()
    .also {
        it.setAnalyzer(cameraExecutor) { imageProxy ->
            val bitmap = imageProxy.toBitmap()
            runInference(bitmap)
            imageProxy.close()  // КРИТИЧНО: иначе CameraX зависнет
        }
    }

imageProxy.close() в блоке finally — не опционально. Если не закрыть ImageProxy, ImageAnalysis перестаёт доставлять новые кадры через несколько секунд. Типичный баг, который обнаруживается только при длительном тестировании.

Числовая точность после конвертации

После конвертации и квантизации обязательно проверяем точность модели на тестовом наборе. FP16 квантизация обычно теряет <1% точности, INT8 — 1–3%. Если потери больше — возможно, квантизационный калибровочный датасет слишком мал или модель чувствительна к определённым слоям.

Для проверки — сравниваем выходы оригинальной PyTorch модели и TFLite на одинаковых входах:

# Тест совпадения выходов
import numpy as np
original_out = pytorch_model(test_input).detach().numpy()
tflite_out = run_tflite(interpreter, test_input)
print(f"Max difference: {np.max(np.abs(original_out - tflite_out))}")
# Норма: < 0.01 для FP16, < 0.05 для INT8

Размещение модели

.tflite файл — в assets/. При первом запуске копируем в filesDir или используем MappedByteBuffer напрямую из assets для zero-copy загрузки:

fun loadModelFile(context: Context, filename: String): MappedByteBuffer {
    val fileDescriptor = context.assets.openFd(filename)
    val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
    return inputStream.channel.map(
        FileChannel.MapMode.READ_ONLY,
        fileDescriptor.startOffset,
        fileDescriptor.declaredLength
    )
}

MappedByteBuffer — OS не копирует файл в RAM при загрузке, а маппирует напрямую. Для больших моделей (50–200 МБ) это существенно.

Процесс

Конвертация из исходного формата → оценка делегатов на целевых устройствах → интеграция с fallback-логикой → тест числовой точности → профилирование через Android Profiler + TFLite Model Benchmark Tool.

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

Базовая интеграция TFLite модели в Android — 1–2 недели. С мультиделегатной логикой, CameraX pipeline, тестированием на парке устройств — 3–5 недель.