將 PyTorch 模型轉換為 LiteRT

AI Edge Torch 是一種程式庫,可讓您將 PyTorch 模型轉換為 .tflite 格式,方便您使用 LiteRT 和 MediaPipe 執行這些模型。 對於要建立可執行模型的行動應用程式開發人員來說,這項功能尤其實用 就能在裝置上進行作業AI Edge Torch 透過初始 GPU 提供廣泛的 CPU 涵蓋範圍 和 NPU 支援

如要開始將 PyTorch 模型轉換為 LiteRT,請使用 Pytorch 轉換工具 快速入門導覽課程。詳情請參閱 AI 邊緣滑桿 GitHub 存放區

如果您特別轉換大型語言模型 (LLM) 建議使用 Generative Torch API。 可處理轉換器專屬的轉換詳細資料,例如模型編寫 量化

轉換工作流程

下列步驟示範 PyTorch 的簡易端對端轉換程序 轉換為 LiteRT

匯入 AI Edge 火箭

首先,請依序匯入 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')