Преобразование моделей PyTorch в TF Lite

AI Edge Torch — это библиотека, позволяющая конвертировать модели PyTorch в формат .tflite , что позволяет запускать эти модели с помощью TensorFlow Lite и MediaPipe. Это особенно полезно для разработчиков, создающих мобильные приложения, которые полностью запускают модели на устройстве. AI Edge Torch обеспечивает широкий охват ЦП с начальной поддержкой графических процессоров и NPU.

Чтобы начать конвертировать модели PyTorch в TF Lite, воспользуйтесь кратким руководством по конвертеру Pytorch . Дополнительную информацию см. в репозитории AI Edge Torch на GitHub .

Если вы специально конвертируете модели большого языка (LLM) или модели на основе преобразователей, используйте Generative Torch API , который обрабатывает детали преобразования, специфичные для преобразователя, такие как создание модели и квантование.

Рабочий процесс преобразования

Следующие шаги демонстрируют простое сквозное преобразование модели PyTorch в TensorFlow Lite.

Импортировать AI Edge Torch

Начните с импорта пакета pip AI Edge Torch ( ai-edge-torch ) вместе с PyTorch.

import ai_edge_torch
import torch

Для этого примера нам также потребуются следующие пакеты:

import numpy
import torchvision

Инициализируйте и преобразуйте модель

Мы конвертируем ResNet18 , популярную модель распознавания изображений.

resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()

Используйте метод convert из библиотеки AI Edge Torch, чтобы преобразовать модель PyTorch.

sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)

Используйте модель

После преобразования модели Pytorch вы можете выполнить выводы с помощью новой преобразованной модели TF Lite.

output = edge_model(*sample_inputs)

Вы можете экспортировать и сохранить преобразованную модель в формате .tflite для дальнейшего использования.

edge_model.export('resnet.tflite')