Конвертация ML-модели в TensorFlow Lite формат для Android

Конвертация ML-модели в TensorFlow Lite формат для Android

TFLite — не просто конвертация весов. Это выбор формата квантизации, оптимизация графа, подбор операционного набора совместимого с целевыми Android-версиями, и проверка того, что числовой результат совпадает с оригиналом. Каждый из этих шагов имеет конкретные грабли.

Пути конвертации

Из TensorFlow SavedModel — прямой путь. Из PyTorch — через ONNX промежуточный формат. Из JAX — через TensorFlow export.

# Путь 1: TF SavedModel → TFLite (наиболее надёжный)
converter = tf.lite.TFLiteConverter.from_saved_model("saved_model_dir/")
tflite_model = converter.convert()

# Путь 2: Keras модель → TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

# Путь 3: PyTorch → ONNX → TF → TFLite
import subprocess
subprocess.run(["python", "-m", "tf2onnx.convert",
    "--onnx", "model.onnx",
    "--output", "model_tf",
    "--opset", "17"])
converter = tf.lite.TFLiteConverter.from_saved_model("model_tf/")

Путь через ONNX вносит дополнительные потенциальные несовместимости — используйте только когда прямой путь недоступен.

Квантизация при конвертации

# FP16 — минимальная деградация, 2× меньше модель, ускорение на GPU delegate
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_fp16 = converter.convert()

# Dynamic INT8 — веса int8, активации float32. Не нужен calibration dataset.
converter2 = tf.lite.TFLiteConverter.from_saved_model("saved_model_dir/")
converter2.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_dynamic_int8 = converter2.convert()

# Full INT8 — и веса, и активации. Требует calibration dataset. Нужен для Hexagon DSP.
def representative_dataset():
    dataset = load_calibration_data()  # 100-500 примеров
    for sample in dataset:
        yield [sample[np.newaxis, :].astype(np.float32)]

converter3 = tf.lite.TFLiteConverter.from_saved_model("saved_model_dir/")
converter3.optimizations = [tf.lite.Optimize.DEFAULT]
converter3.representative_dataset = representative_dataset
converter3.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter3.inference_input_type = tf.uint8   # или tf.int8
converter3.inference_output_type = tf.uint8
tflite_full_int8 = converter3.convert()

Full INT8 с inference_input_type = tf.uint8 — входные данные передаются как uint8 (0–255), не нужно нормализовать в float32 в Java/Kotlin. Это убирает один шаг предобработки, но требует аккуратного согласования с quantization parameters модели.

Неподдерживаемые операции

Не все TF/PyTorch операции есть в TFLite builtin ops. Проверка:

# Какие операции в модели не поддерживаются TFLite
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS  # fallback на TF операции
]

SELECT_TF_OPS подключает подмножество TF операций — это увеличивает размер бинарника TFLite runtme (~5 МБ) и замедляет некоторые операции. Лучше переписать модель чтобы обойтись без SELECT_TF_OPS — это даёт совместимость с NNAPI и Hexagon.

Кастомная операция через C++:

// Регистрация кастомного оператора
static TfLiteRegistration* GetMyCustomOpRegistration() {
    static TfLiteRegistration reg = {
        nullptr, nullptr,
        [](TfLiteContext* ctx, TfLiteNode* node) -> TfLiteStatus {
            // Реализация инференса
            return kTfLiteOk;
        },
        nullptr, "MyCustomOp", 1
    };
    return ®
}

// В Android NDK коде:
interpreter.AddCustomOp("MyCustomOp", GetMyCustomOpRegistration, 1);

В Kotlin через JNI. Это нетривиально, но иногда единственный путь.

Верификация числовой точности

# Сравнение PyTorch/TF и TFLite выходов на одинаковых входах
import numpy as np

# TF оригинал
tf_output = tf_model(test_input).numpy()

# TFLite
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
tflite_output = interpreter.get_tensor(output_details[0]['index'])

print(f"Max abs diff: {np.max(np.abs(tf_output - tflite_output))}")
print(f"MSE: {np.mean((tf_output - tflite_output)**2)}")
# FP32: < 1e-5, FP16: < 1e-2, INT8: < 0.05

Если разница больше нормы — проблема в нормализации входных данных, неправильных quantization parameters или в операции, для которой TFLite использует другой алгоритм.

Особенности для детекторов объектов

YOLO, SSD, EfficientDet — содержат NMS (Non-Maximum Suppression) постпроцессинг. TFLite не умеет NMS встроенно (в отличие от Core ML Detection Output). Варианты:

1. Убрать NMS из модели, реализовать в Java/Kotlin после инференса
2. Использовать TFLite Task Library — она содержит готовый ObjectDetection API с NMS

// TFLite Task Library: ObjectDetector (включает NMS)
val options = ObjectDetector.ObjectDetectorOptions.builder()
    .setScoreThreshold(0.5f)
    .setMaxResults(20)
    .build()

val detector = ObjectDetector.createFromFileAndOptions(context, "detector.tflite", options)

val image = TensorImage.fromBitmap(inputBitmap)
val results: List<Detection> = detector.detect(image)

for (detection in results) {
    val box = detection.boundingBox  // RectF
    val label = detection.categories.first().label
    val score = detection.categories.first().score
}

Task Library поддерживает только определённые model signatures — модель должна следовать формату TFLite Model Metadata.

Model Metadata: важно для Interpreter + Task Library

from tflite_support.metadata_writers import image_classifier
from tflite_support.metadata_writers import writer_utils

# Создаём метаданные для классификатора
writer = image_classifier.MetadataWriter.create_for_inference(
    writer_utils.load_file("model.tflite"),
    input_norm_mean=[0.0],
    input_norm_std=[255.0],
    labels_file_paths=["labels.txt"]
)

tflite_with_metadata = writer.populate()
writer_utils.save_file(tflite_with_metadata, "model_with_metadata.tflite")

Без метаданных TFLite Task Library работает хуже — нет автоматической нормализации, нет маппинга выходов. С метаданными — всё обрабатывается автоматически.

Бенчмарк на устройствах

# ADB: запускаем TFLite Benchmark Tool напрямую на устройстве
adb push model.tflite /data/local/tmp/
adb shell /data/local/tmp/benchmark_model \
    --graph=/data/local/tmp/model.tflite \
    --use_gpu=true \
    --num_threads=4 \
    --num_runs=50
# Вывод: средняя задержка, min/max, warmup время

TFLite Benchmark Tool — официальный инструмент от Google, даёт честные цифры без JVM overhead и Android UI. Использовать для сравнения делегатов (GPU vs NNAPI vs CPU).

Процесс

Выбор пути конвертации → конвертация с нужным уровнем квантизации → верификация точности → добавление метаданных → тест на парке устройств через Benchmark Tool → интеграция в Android-приложение.

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

Прямая конвертация TF/Keras модели с верификацией — 3–7 дней. Конвертация через ONNX, кастомные операции, добавление метаданных, полное тестирование — 2–4 недели.