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')