Конвертация ML-модели в Core ML формат для iOS

Конвертация ML-модели в Core ML формат для iOS

Получить .mlpackage из PyTorch или TensorFlow модели — задача с конкретными шагами, где каждый шаг может сломаться по специфической причине. coremltools конвертирует не всё и не всегда правильно. Разберём реальный процесс конвертации, включая то, что идёт не так.

Подготовка модели к конвертации

Перед конвертацией модель должна быть в eval-режиме с фиксированными весами. torch.jit.trace требует примера входных данных — он записывает граф выполнения для конкретного shape:

import torch
import coremltools as ct

model = MyModel()
model.load_state_dict(torch.load("weights.pth", map_location="cpu"))
model.eval()

# trace — фиксирует граф для конкретного shape
example_input = torch.zeros(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)

# Для моделей с условными ветками (if/else зависящими от данных) — script:
# scripted_model = torch.jit.script(model)  # анализирует код, не данные

torch.jit.trace не работает с динамическими ветками — если в forward() есть if x.shape[1] > 512: ..., trace запомнит только одну ветку. В таком случае нужен torch.jit.script, который анализирует весь код статически.

Конвертация через coremltools

mlmodel = ct.convert(
    traced_model,
    inputs=[ct.ImageType(
        name="input",
        shape=ct.Shape(shape=(1, 3, 224, 224)),
        color_layout=ct.colorlayout.RGB,
        # Нормализация встроена в модель — не нужно делать в Swift
        bias=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
        scale=1/(255.0 * 0.229)
    )],
    outputs=[ct.TensorType(name="logits")],
    compute_precision=ct.precision.FLOAT16,
    minimum_deployment_target=ct.target.iOS16,
    convert_to="mlprogram"  # новый формат .mlpackage
)

# Добавляем метаданные
mlmodel.short_description = "Image classifier"
mlmodel.input_description["input"] = "RGB image 224x224"
mlmodel.output_description["logits"] = "Class probabilities"

mlmodel.save("MyModel.mlpackage")

convert_to="mlprogram" vs "neuralnetwork": mlprogram — новый IR, поддерживает FP16 ANE-операции, требует iOS 15+. neuralnetwork — старый формат, iOS 12+. Для новых проектов — mlprogram.

Частые ошибки конвертации

ValueError: TorchScript conversion failed — обычно из-за нестандартных операций. Проверяем через:

# Какие операции не поддерживаются:
print(ct.converters.mil.frontend.torch.ops.PYTORCH_OPS_REGISTRY.keys())
# Если операции нет — нужна кастомная реализация

RuntimeError: Only tuple/list outputs are supported — модель возвращает dict. Нужна обёртка:

class ModelWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, x):
        out = self.model(x)
        if isinstance(out, dict):
            return out["logits"]  # извлекаем нужный ключ
        return out

Unsupported op: aten::einsum — einsum в PyTorch. Нужно переписать через стандартные matmul/bmm, или добавить custom op.

ONNX как промежуточный шаг — когда прямая конвертация не работает:

# PyTorch → ONNX → Core ML
torch.onnx.export(model, example_input, "model.onnx", opset_version=17)
mlmodel = ct.converters.onnx.convert(
    model="model.onnx",
    minimum_ios_deployment_target="15.0"
)

Конвертация через ONNX иногда решает проблемы с неподдерживаемыми PyTorch операциями.

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

import numpy as np
import PIL.Image

# Загружаем тестовое изображение
img = PIL.Image.open("test.jpg").resize((224, 224))

# Оригинальный PyTorch вывод
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
tensor = transform(img).unsqueeze(0)
with torch.no_grad():
    pytorch_out = model(tensor).numpy()

# Core ML вывод
coreml_out = mlmodel.predict({"input": img})["logits"]

# Сравниваем
max_diff = np.max(np.abs(pytorch_out - coreml_out))
print(f"Max difference: {max_diff}")
# Норма для FP16: < 0.01
# Если больше 0.05 — что-то пошло не так при конвертации

Если разница большая — проверяем нормализацию: возможно, bias и scale в ct.ImageType не соответствуют нормализации в pytorch transform.

Переменные размеры входа

Для моделей, работающих с изображениями разного размера:

# Диапазон размеров
flexible_shape = ct.Shape(
    shape=(1, 3, ct.RangeDim(min_val=64, max_val=1024), ct.RangeDim(min_val=64, max_val=1024))
)

# Или набор конкретных размеров
enumerated_shapes = ct.EnumeratedShapes(
    shapes=[
        ct.Shape(shape=(1, 3, 224, 224)),
        ct.Shape(shape=(1, 3, 384, 384)),
        ct.Shape(shape=(1, 3, 512, 512)),
    ]
)

mlmodel = ct.convert(traced_model, inputs=[ct.TensorType(name="input", shape=enumerated_shapes)])

EnumeratedShapes позволяет Core ML заранее оптимизировать граф для каждого из указанных размеров. RangeDim — более гибкий, но оптимизация хуже.

Кастомные операции

Если модель содержит операцию, которую coremltools не знает — можно добавить кастомный слой на Swift/Objective-C:

# Python: регистрация кастомной операции
@ct.converters.mil.register_torch_op()
def my_custom_op(context, node):
    # Реализация конвертации через MIL операции
    x = context[node.inputs[0]]
    result = mb.custom(params={"...": "..."}, inputs={"x": x}, ...)
    context.add(result)

// Swift: реализация кастомного слоя
import CoreML

@objc(MyCustomLayer)
class MyCustomLayer: NSObject, MLCustomLayer {
    required init(parameters: [String: Any]) throws { }

    func setWeightData(_ weights: [Data]) throws { }
    func outputShapes(forInputShapes inputShapes: [[NSNumber]]) throws -> [[NSNumber]] { ... }
    func evaluate(inputs: [MLMultiArray], outputs: [MLMultiArray]) throws { ... }
}

Кастомные слои всегда выполняются на CPU, без ANE/GPU ускорения. Если custom op — узкое место инференса, лучше переписать модель на стандартных операциях.

Интеграция .mlpackage в Xcode

Файл .mlpackage добавляется в проект — Xcode автоматически генерирует Swift-класс с typed API. Имя файла MyModel.mlpackage → класс MyModel с методом prediction(input:).

Размер .mlpackage в bundle влияет на время загрузки в App Store и размер скачиваемого приложения. Для больших моделей (>50 МБ) — размещать вне bundle, скачивать при первом запуске через URLSession и хранить в Application Support.

Процесс

Аудит модели на совместимость → конвертация с подбором параметров → верификация числовой точности → тестирование на целевых iOS-версиях → профилирование в Xcode Core ML Instrument.

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

Стандартная модель без нестандартных операций — 3–7 дней. Модель со сложным графом, кастомными операциями, переменными размерами — 2–4 недели.