AI Edge Torch — это библиотека, позволяющая конвертировать модели PyTorch в формат .tflite , что позволяет запускать эти модели с помощью LiteRT и MediaPipe. Это особенно полезно для разработчиков, создающих мобильные приложения, которые полностью запускают модели на устройстве. AI Edge Torch обеспечивает широкий охват ЦП с начальной поддержкой графических процессоров и NPU.
Чтобы приступить к преобразованию моделей PyTorch в LiteRT, воспользуйтесь кратким руководством по конвертеру Pytorch . Дополнительную информацию см. в репозитории AI Edge Torch на GitHub .
Если вы специально конвертируете модели большого языка (LLM) или модели на основе преобразователей, используйте Generative Torch API , который обрабатывает детали преобразования, специфичные для преобразователя, такие как создание модели и квантование.
Рабочий процесс преобразования
Следующие шаги демонстрируют простое сквозное преобразование модели PyTorch в LiteRT.
Импортировать 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 вы можете выполнить выводы с помощью новой преобразованной модели LiteRT.
output = edge_model(*sample_inputs)
Вы можете экспортировать и сохранить преобразованную модель в формате .tflite для дальнейшего использования.
edge_model.export('resnet.tflite')