將 PyTorch 模型轉換為 TF Lite

AI Edge Torch 程式庫可讓您將 PyTorch 模型轉換為 .tflite 格式，以便透過 TensorFlow Lite 和 MediaPipe 執行這些模型。如果開發人員要建立能在裝置端完整執行模型的行動應用程式，這項功能就特別實用。AI Edge Torch 提供廣泛的 CPU 涵蓋範圍 包括初始 GPU 和 NPU 支援

如要開始將 PyTorch 模型轉換為 TF Lite，請使用 Pytorch 轉換工具快速入門導覽課程。詳情請參閱 AI Edge Torch GitHub 存放區。

如果您要轉換大型語言模型 (LLM) 或轉換型模型，請使用 Generative Torch API。這個 API 可處理轉換器專屬的轉換詳細資料，例如模型編寫和量化。

轉換工作流程

下列步驟示範如何將 PyTorch 模型的端對端端對端轉換為 TensorFlow Lite。

匯入 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 模型後，您可以使用新的轉換後的 TF Lite 模型執行推論。

output = edge_model(*sample_inputs)

您可以匯出轉換後的模型，並儲存為 .tflite 格式，以供日後使用。

edge_model.export('resnet.tflite')