將 PyTorch 模型轉換為 LiteRT

AI Edge Torch 是一個程式庫，可將 PyTorch 模型轉換為 .tflite 格式，讓您使用 LiteRT 和 MediaPipe 執行這些模型。這對於開發人員而言特別實用，因為他們開發的行動應用程式可完全在裝置上執行模型。AI Edge Torch 提供廣泛的 CPU 涵蓋範圍，並提供初始 GPU 和 NPU 支援。

如要開始將 PyTorch 模型轉換為 LiteRT，請使用 Pytorch 轉換工具 quickstart。詳情請參閱 AI Edge Torch GitHub 存放區。

如果您要特別轉換大型語言模型 (LLM) 或以轉換器為基礎的模型，請使用 Generative Torch API，該 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')