Как конвертировать сегментацию pytorch в CoreML

Вопрос или проблема

import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import coremltools as ct

# Загружаем предобученную модель DeepLabV3 с основной архитектурой MobileNetV3
model = torch.hub.load('pytorch/vision:v0.19.0', 'deeplabv3_mobilenet_v3_large', pretrained=True)
model.eval()

# Класс-обертка
class Deeplabv3Wrapper(nn.Module):
    def __init__(self, model):
        super(Deeplabv3Wrapper, self).__init__()
        self.model = model

    def forward(self, x):
        outputs = self.model(x)
        return outputs['out']  # Возвращаем тензор 'out' для маски сегментации

# Определяем преобразования для предварительной обработки
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # значения из документации pytorch
])

dummy_image = Image.new("RGB", (513, 513), color="white")
dummy_tensor = preprocess(dummy_image).unsqueeze(0)  # Добавляем размерность для батча

# Оборачиваем и трассируем модель
wrapped_model = Deeplabv3Wrapper(model)

# Трассируем модель с помощью TorchScript
traced_model = torch.jit.trace(wrapped_model, dummy_tensor)
print("Трассировка модели завершена")

# Сохраняем трассированную модель
torch.jit.save(traced_model, "traced_deeplabv3.pt") # ЭТО МОЖНО открыть в netron.app
print("Модель успешно сохранена в формате pt")


traced_model.eval()

example_input = torch.randn(1, 3, 513, 513)


model_from_trace = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name="input_image", shape=example_input.shape, scale=1/255.0, bias=[-0.485, -0.456, -0.406])],
    outputs=[ct.TensorType(name="output")]
)

print("Конвертация модели завершена")

model_from_trace.save("DeepLabV3_MobileNet.mlpackage") # ЭТУ МОДЕЛЬ НЕЛЬЗЯ открыть в netron.app - Ошибка: "Файл не содержит данных"

print("Модель успешно сохранена в формате mlpackage")

Ответ или решение

Пошаговое руководство по конвертации сегментации PyTorch в CoreML

Конвертация моделей глубокого обучения из PyTorch в CoreML станет полезной задачей для многих разработчиков, стремящихся использовать мощь компьютерного зрения в своих приложениях Apple. В этой статье мы рассмотрим, как успешно выполнить эту конвертацию на примере модели DeepLabV3 с использованием MobileNetV3 в качестве основы. Мы также разберём, как избежать распространённых ошибок, связанных с работой с форматами моделей и преобразованиями.

Шаг 1: Установка необходимых библиотек

Перед тем как начать процесс конвертации, убедитесь, что у вас установлены необходимые зависимости. Ваша версия PyTorch должна быть совместима с CoreML. На данный момент рекомендуется использовать следующие версии:

pip install "torch==2.4" "torchvision==0.19" coremltools

Шаг 2: Загрузка модели DeepLabV3

Используйте следующий код для загрузки модели DeepLabV3 с предобученной MobileNetV3:

import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image

# Загрузка модели DeepLabV3
model = torch.hub.load('pytorch/vision:v0.19.0', 'deeplabv3_mobilenet_v3_large', pretrained=True)
model.eval()

Шаг 3: Обёртка модели

Создайте класс-обёртку для вашей модели, который упростит процесс её трассировки:

class Deeplabv3Wrapper(nn.Module):
    def __init__(self, model):
        super(Deelabv3Wrapper, self).__init__()
        self.model = model

    def forward(self, x):
        outputs = self.model(x)
        return outputs['out']  # Возвращаем выходной тензор для сегментации

Шаг 4: Предобработка данных

Необходимо подготовить ваши входные данные с использованием стандартных преобразований:

preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # Нормализация в соответствии с документацией PyTorch
])

Шаг 5: Трассировка модели

Создайте тестовый образец и выполните трассировку модели с использованием TorchScript:

dummy_image = Image.new("RGB", (513, 513), color="white")
dummy_tensor = preprocess(dummy_image).unsqueeze(0)  # Добавляем размерность пакета

wrapped_model = Deeplabv3Wrapper(model)

traced_model = torch.jit.trace(wrapped_model, dummy_tensor)
print("Трассировка модели завершена")

torch.jit.save(traced_model, "traced_deeplabv3.pt")  # Сохранение модели
print("Модель успешно сохранена как .pt")

Шаг 6: Конвертация в формат CoreML

Теперь мы можем конвертировать трассированную модель в формат CoreML. Обратите внимание на необходимость правильно указать входные данные:

import coremltools as ct

example_input = torch.randn(1, 3, 513, 513)

model_from_trace = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name="input_image", shape=example_input.shape, scale=1/255.0, bias=[-0.485, -0.456, -0.406])],
    outputs=[ct.TensorType(name="output")]
)

print("Конвертация модели завершена")

model_from_trace.save("DeepLabV3_MobileNet.mlpackage")  # Сохранение в формате mlpackage
print("Модель успешно сохранена как .mlpackage")

Шаг 7: Устранение ошибок "Нет контента"

Если вы столкнулись с ошибкой "File has no content" при попытке открывать файл .mlpackage в Netron, проверьте следующее:

1. Совместимость версий: Убедитесь, что используемые вами версии библиотек (torch, torchvision и coremltools) корректны и совместимы друг с другом.

2. Также попробуйте вместо ct.TensorType(name="output") использовать другие форматы на выходе, так как иногда CoreML требует явного указания формата выходного тензора.

3. Переустановите библиотеки: Иногда вмешательство других пакетов может нарушить функциональность. Переустановка может решить эту проблему.

Заключение

Следуя этому пошаговому руководству, вы сможете легко конвертировать модели сегментации из PyTorch в CoreML для использования в приложениях Apple. Убедитесь, что вы правильно настраиваете входные данные и форматы выходов, чтобы избежать распространённых проблем с несовместимостью. Если у вас возникли вопросы или сложности в процессе, не стесняйтесь использовать сообщества разработчиков для получения помощи.