将 PyTorch 模型转换为 LiteRT

AI Edge Torch 是一个库,可让您将 PyTorch 模型转换为 .tflite 这样你就可以使用 LiteRT 和 MediaPipe 运行这些模型了。 对于要开发运行模型的移动应用的开发者而言,这尤其有用 完全在设备上完成操作AI Edge Torch 提供广泛的 CPU 覆盖范围,初始 GPU 和 NPU 支持。

如需开始将 PyTorch 模型转换为 LiteRT,请使用 Pytorch 转换器 快速入门。如需了解详情,请参阅 AI Edge Torch GitHub 代码库

如果您要专门将大语言模型 (LLM) 或 基于 Transformer 的模型,请使用 Generative Torch API, 用于处理特定于转换器的转换详情,例如模型编写和 量化。

转化工作流程

以下步骤演示了 PyTorch 的简单端到端转换 转换为 LiteRT。

导入 AI Edge Torch

首先导入 AI Edge Torch (ai-edge-torch) pip 软件包以及 PyTorch。

import ai_edge_torch
import torch

在此示例中,我们还需要以下软件包:

import numpy
import torchvision

初始化并转换模型

我们会将 ResNet18、 一种流行的图片识别模型。

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

使用 AI Edge Torch 库中的 convert 方法转换 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')